| 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) | |
| # 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 | |
| # 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 | |
| # 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() | |