| import torch |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
| class PythiaChatHandler: |
| def __init__(self, model_name="togethercomputer/Pythia-Chat-Base-7B"): |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
| self.model = AutoModelForCausalLM.from_pretrained(model_name).to(self.device) |
|
|
| def __call__(self, input_text): |
| input_ids = self.tokenizer.encode(input_text, return_tensors="pt").to(self.device) |
| output_ids = self.model.generate(input_ids).to("cpu") |
| response_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True) |
| return response_text |
|
|
