blog / model.py
Mo-alaa's picture
Update model.py
7ab159a verified
raw
history blame
324 Bytes
from transformers import AutoModelForCausalLM, AutoTokenizer
def create_model():
tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-zephyr-3b')
model = AutoModelForCausalLM.from_pretrained(
'stabilityai/stablelm-zephyr-3b',
trust_remote_code=True,
device_map="auto"
)
return model,tokenizer