Upload handler.py

handler.py

@@ -0,0 +1,14 @@
+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