Spaces:
Sleeping
Sleeping
File size: 744 Bytes
a9292a1 5cca49a a9292a1 e79ab5a 65c9e0f a9292a1 65c9e0f a9292a1 65c9e0f a9292a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
"""
File: model_llm.py
Description: Load a Large Language Model (LLM)
Author: Didier Guillevic
Date: 2024-03-16
"""
import torch
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import TextIteratorStreamer
import os
model_name = "mistralai/Mistral-7B-Instruct-v0.3"
auth_token = os.environ.get("HF_TOKEN")
# Auto-regressive model for language completion: padding left
tokenizer = AutoTokenizer.from_pretrained(
model_name,
padding_side="left",
use_auth_token=auth_token
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
torch_dtype=torch.float16,
low_cpu_mem_usage=True,
use_auth_token=auth_token
)
model = torch.compile(model)
|