Update app.py

app.py

@@ -1,85 +1,89 @@
+
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
-# Load the Salesforce/codet5-base model and tokenizer
-# We are using the 'Salesforce/codet5-base' model, which is a pre-trained model for code-related tasks.
-# The AutoTokenizer and AutoModelForSeq2SeqLM classes from the Transformers library are used to load the model and tokenizer.
-model_name = "Salesforce/codet5-base"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# 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.
+        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)
+
+    # 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)
+
+# 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):
+        self.history = []
+
+    # 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.
+        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 create_gradio_interface():
+    # Create an instance of the CodeGenerator class.
+    code_generator = CodeGenerator()
+    # Create an instance of the ChatHandler class.
+    chat_handler = ChatHandler()
+
+    # 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")
 
-# Function to generate code
-# This function takes a prompt (code-related query) as input and generates code based on that prompt.
-# It uses the loaded model and tokenizer to encode the input, generate the output, and then decode the generated text.
-def generate_code(prompt, max_length=100):
-    # Encode the input prompt using the tokenizer
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    
-    # Generate the output using the model
-    # The `model.generate()` function is used to generate the output sequence based on the input.
-    # The `max_length` parameter sets the maximum length of the generated sequence.
-    # The `num_return_sequences` parameter specifies the number of output sequences to be generated (in this case, 1).
-    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1)
-    
-    # Decode the generated output to get the actual code
-    # The `tokenizer.decode()` function is used to convert the output token IDs back to readable text.
-    # The `skip_special_tokens=True` argument ensures that any special tokens (e.g., start/end of sequence tokens) are removed from the output.
-    generated_code = tokenizer.decode(output[0], skip_special_tokens=True)
-    
-    # Return the generated code
-    return generated_code
+            # 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")
 
-# Function to handle chat interaction
-# This function is responsible for managing the chat interaction between the user and the system.
-# It takes the user's message and the chat history as input, and returns the system's response and the updated chat history.
-def chat_interaction(message, history):
-    # Initialize the chat history if it's not provided
-    history = history or []
-    
-    # Generate the response using the `generate_code` function
-    response = generate_code(message)
-    
-    # Update the chat history by appending the user's message and the system's response
-    history.append((message, response))
-    
-    # Return the empty message (to clear the input field) and the updated chat history
-    return "", history
+        # 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=[])
 
-# Create the Gradio interface
-# The Gradio library is used to create an interactive web interface for the chat application.
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    # Add a Markdown title for the interface
-    gr.Markdown("# S-Dreamer Salesforce/codet5-base Chat Interface")
-    
-    # Create a row with two columns
-    with gr.Row():
-        # Left column for the chat area
-        with gr.Column(scale=3):
-            # Add a chatbot component to display the chat history
-            chatbot = gr.Chatbot(height=400)
-            # Add a text input field for the user to enter messages
-            message = gr.Textbox(label="Enter your code-related query", placeholder="Type your message here...")
-            # Add a submit button
-            submit_button = gr.Button("Submit")
-        
-        # Right column for the feature list
-        with gr.Column(scale=1):
-            # Add Markdown sections for the features
-            gr.Markdown("## Features")
-            gr.Markdown("- Code generation")
-            gr.Markdown("- Code completion")
-            gr.Markdown("- Code explanation")
-            gr.Markdown("- Error correction")
-            
-            # Add a clear button to reset the chat
-            clear_button = gr.Button("Clear Chat")
-    
-    # Connect the submit button to the `chat_interaction` function
-    submit_button.click(chat_interaction, inputs=[message, chatbot], outputs=[message, chatbot])
-    
-    # Connect the clear button to a lambda function that clears the chat
-    clear_button.click(lambda: None, outputs=[chatbot], inputs=[])
+    # Launch the Gradio interface.
+    demo.launch()
 
-# 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()