| 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) | |