gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

move_code

Browse Source
This commit is contained in:
tackhwa
2024-10-23 15:25:46 +08:00
parent 63c0283514
commit fd30ae4e45
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
lightrag/llm.py
View File

@@ -207,6 +207,8 @@ async def bedrock_complete_if_cache(
def initialize_hf_model(model_name): def initialize_hf_model(model_name):
hf_tokenizer = AutoTokenizer.from_pretrained(model_name, device_map="auto", trust_remote_code=True) hf_tokenizer = AutoTokenizer.from_pretrained(model_name, device_map="auto", trust_remote_code=True)
hf_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", trust_remote_code=True) hf_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", trust_remote_code=True)
if hf_tokenizer.pad_token is None:
hf_tokenizer.pad_token = hf_tokenizer.eos_token
return hf_model, hf_tokenizer return hf_model, hf_tokenizer
@@ -216,9 +218,6 @@ async def hf_model_if_cache(
) -> str: ) -> str:
model_name = model model_name = model
hf_model, hf_tokenizer = initialize_hf_model(model_name) hf_model, hf_tokenizer = initialize_hf_model(model_name)
if hf_tokenizer.pad_token is None:
# print("use eos token")
hf_tokenizer.pad_token = hf_tokenizer.eos_token
hashing_kv: BaseKVStorage = kwargs.pop("hashing_kv", None) hashing_kv: BaseKVStorage = kwargs.pop("hashing_kv", None)
messages = [] messages = []
if system_prompt: if system_prompt: