model.py

from transformers import AutoModelForCausalLM, AutoTokenizer

@torch.no_grad()
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