Create chatbot_model.py

chatbot_model.py

@@ -0,0 +1,13 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+class ChatbotModel:
+    def __init__(self, model_name: str):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name)
+
+    def get_response(self, user_message: str) -> str:
+        # Tokenize input and generate response
+        inputs = self.tokenizer.encode(user_message + self.tokenizer.eos_token, return_tensors="pt")
+        outputs = self.model.generate(inputs, max_length=150, num_return_sequences=1)
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response.strip()