Create inference.py

inference.py

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