weisong
/
llama-2-7b-chat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update model.py
Build / Build-Deploy-Actions (push) Has been cancelled Details

This commit is contained in:
weisong 2024-07-24 15:55:39 +08:00
parent ae133daced
commit e36785d628
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
model.py
View File

@ -7,7 +7,7 @@ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIt
model_id = 'meta-llama/Llama-2-7b-chat-hf'
if torch.cuda.is_available():
config = AutoConfig.from_pretrained(model_id, use_auth_token ='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')
config = AutoConfig.from_pretrained(model_id, use_auth_token ='hf_*****************************')
config.pretraining_tp = 1
model = AutoModelForCausalLM.from_pretrained(
model_id,
@ -15,11 +15,11 @@ if torch.cuda.is_available():
torch_dtype=torch.float16,
load_in_4bit=True,
device_map='auto',
use_auth_token ='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'
use_auth_token ='hf_******************************************'
)
else:
model = None
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token ='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')
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token ='hf_******************************')
def get_prompt(message: str, chat_history: list[tuple[str, str]],