Update app.py

app.py

@@ -3,15 +3,15 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
 class CodeGenerator:
-    def __init__(self, model_name="Salesforce/codet5-base", device=None):
+    def __init__(self, model_name="Salesforce/codet5-base"):
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-        if device:
-            self.model = self.model.to(device)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
 
     def generate_code(self, prompt, max_length=100):
         try:
-            input_ids = self.tokenizer.encode(prompt, return_tensors="pt")
+            input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
             output = self.model.generate(input_ids, max_length=max_length, num_return_sequences=1)
             return self.tokenizer.decode(output[0], skip_special_tokens=True)
         except Exception as e:
@@ -20,13 +20,14 @@ class CodeGenerator:
 class ChatHandler:
     def __init__(self, code_generator):
         self.history = []
-        self.code_generator = code_generator  # Store the generator reference
+        self.code_generator = code_generator
 
     def handle_message(self, message):
         if not message.strip():
             return "", self.history
         response = self.code_generator.generate_code(message)
-        self.history.append((message, response))
+        self.history.append({"role": "user", "content": message})
+        self.history.append({"role": "assistant", "content": response})
         return "", self.history
 
     def clear_history(self):
@@ -34,8 +35,7 @@ class ChatHandler:
         return []
 
 def create_gradio_interface():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    code_generator = CodeGenerator(device=device)
+    code_generator = CodeGenerator()
     chat_handler = ChatHandler(code_generator)
 
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -43,7 +43,7 @@ def create_gradio_interface():
 
         with gr.Row():
             with gr.Column(scale=3):
-                chatbot = gr.Chatbot(height=400)
+                chatbot = gr.Chatbot(type="messages", height=400)
                 message_input = gr.Textbox(label="Enter your code-related query", placeholder="Type your message here...")
                 submit_button = gr.Button("Submit")
 
@@ -60,4 +60,4 @@ def create_gradio_interface():
     demo.launch()
 
 if __name__ == "__main__":
-    create_gradio_interface()
+    create_gradio_interface()