Create chatbot.py

modules/chatbot.py

@@ -0,0 +1,21 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+class Llama2Chatbot:
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
+        self.model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+
+    def generate_response(self, prompt, max_length=100):
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        outputs = self.model.generate(**inputs, max_length=max_length)
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+
+def initialize_chatbot():
+    return Llama2Chatbot()
+
+def get_chatbot_response(chatbot, prompt):
+    return chatbot.generate_response(prompt)