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