File size: 792 Bytes
943e809 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
class Model:
def __init__(self):
device = "cuda" if torch.cuda.is_available() else "cpu"
self.tokenizer = AutoTokenizer.from_pretrained(".")
self.model = AutoModelForCausalLM.from_pretrained(
".",
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
device_map="auto"
)
def __call__(self, prompt: str):
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
output_ids = self.model.generate(
**inputs,
max_new_tokens=128,
temperature=0.7,
do_sample=True
)
return self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
|