Update app.py

app.py

@@ -1,89 +1,63 @@
-
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
-# Define a class named `CodeGenerator` that will be responsible for generating code based on a given prompt.
 class CodeGenerator:
-    # The constructor initializes the CodeGenerator object with a pre-trained model name.
-    # The default model name is "Salesforce/codet5-base".
-    def __init__(self, model_name="Salesforce/codet5-base"):
-        # Load the pre-trained tokenizer from the specified model name.
+    def __init__(self, model_name="Salesforce/codet5-base", device=None):
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        # Load the pre-trained sequence-to-sequence language model from the specified model name.
         self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        if device:
+            self.model = self.model.to(device)
 
-    # This method generates code based on the given prompt.
-    # The method takes two parameters: `prompt` (the input text) and `max_length` (the maximum length of the generated code).
     def generate_code(self, prompt, max_length=100):
-        # Encode the prompt into input IDs that the model can understand.
-        input_ids = self.tokenizer.encode(prompt, return_tensors="pt")
-        # Generate the output sequence using the pre-trained model.
-        # The `generate` method takes the input IDs, the maximum length of the output, and the number of output sequences to return (in this case, 1).
-        output = self.model.generate(input_ids, max_length=max_length, num_return_sequences=1)
-        # Decode the output sequence and return the generated code.
-        return self.tokenizer.decode(output[0], skip_special_tokens=True)
+        try:
+            input_ids = self.tokenizer.encode(prompt, return_tensors="pt")
+            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:
+            return f"Error generating code: {str(e)}"
 
-# Define a class named `ChatHandler` that will be responsible for managing the chat history.
 class ChatHandler:
-    # The constructor initializes the ChatHandler object with an empty chat history.
-    def __init__(self):
+    def __init__(self, code_generator):
         self.history = []
+        self.code_generator = code_generator  # Store the generator reference
 
-    # This method handles incoming messages and generates responses using the provided CodeGenerator.
-    # The method takes two parameters: `message` (the user's input message) and `code_generator` (an instance of the CodeGenerator class).
-    def handle_message(self, message, code_generator):
-        # Generate the response using the provided CodeGenerator.
-        response = code_generator.generate_code(message)
-        # Append the message-response pair to the chat history.
+    def handle_message(self, message):
+        if not message.strip():
+            return "", self.history
+        response = self.code_generator.generate_code(message)
         self.history.append((message, response))
-        # Return the empty message input and the updated chat history.
         return "", self.history
 
-# Define a function named `create_gradio_interface` that creates a Gradio interface for the chat application.
+    def clear_history(self):
+        self.history = []
+        return []
+
 def create_gradio_interface():
-    # Create an instance of the CodeGenerator class.
-    code_generator = CodeGenerator()
-    # Create an instance of the ChatHandler class.
-    chat_handler = ChatHandler()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    code_generator = CodeGenerator(device=device)
+    chat_handler = ChatHandler(code_generator)
 
-    # Create a Gradio Blocks interface with a soft theme.
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        # Display a Markdown title for the chat interface.
         gr.Markdown("# S-Dreamer Salesforce/codet5-base Chat Interface")
 
-        # Create a row with two columns.
         with gr.Row():
-            # The first column will contain the chat interface.
             with gr.Column(scale=3):
-                # Create a chatbot component to display the chat history.
                 chatbot = gr.Chatbot(height=400)
-                # Create a textbox for the user to input their message.
                 message_input = gr.Textbox(label="Enter your code-related query", placeholder="Type your message here...")
-                # Create a submit button to send the message.
                 submit_button = gr.Button("Submit")
 
-            # The second column will contain the features.
             with gr.Column(scale=1):
-                # Display a Markdown title for the features section.
                 gr.Markdown("## Features")
-                # Define a list of features.
                 features = ["Code generation", "Code completion", "Code explanation", "Error correction"]
-                # Display each feature as a Markdown list item.
                 for feature in features:
                     gr.Markdown(f"- {feature}")
-                # Create a button to clear the chat history.
                 clear_button = gr.Button("Clear Chat")
 
-        # Connect the submit button to the `handle_message` method of the ChatHandler.
-        submit_button.click(chat_handler.handle_message, inputs=[message_input], outputs=[message_input, chatbot])
-        # Connect the clear button to a function that clears the chat history.
-        clear_button.click(lambda: None, outputs=[chatbot], inputs=[])
+        submit_button.click(chat_handler.handle_message, inputs=message_input, outputs=[message_input, chatbot])
+        clear_button.click(lambda: (None, chat_handler.clear_history()), inputs=[], outputs=[message_input, chatbot])
 
-    # Launch the Gradio interface.
     demo.launch()
 
-# This is the entry point of the application.
 if __name__ == "__main__":
-    # Call the `create_gradio_interface` function to start the chat application.
     create_gradio_interface()