Update 'model.py'
Build-Deploy-Actions
Details
Build-Deploy-Actions
Details
This commit is contained in:
parent
7cfb612cab
commit
b9a1ddb80f
6
model.py
6
model.py
|
@ -7,7 +7,7 @@ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIt
|
|||
model_id = 'meta-llama/Llama-2-13b-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]],
|
||||
|
|
Loading…
Reference in New Issue