import gradio as gr | |
from transformers import pipeline | |
# Initialize the StarCoder model | |
# Replace 'models/bigcode/starcoder' with the actual model path if different | |
code_generator = pipeline('text-generation', model='google/codegemma-2b') | |
def generate_code(prompt): | |
#""" | |
#Generates code based on the given prompt. | |
#Parameters: | |
#- prompt (str): The input text to generate code from. | |
#Returns: | |
#- str: The generated code. | |
#""" | |
generated_code = code_generator(prompt, max_length=100, do_sample=True, temperature=0.7) | |
return generated_code[0]['generated_text'] | |
# Create the Gradio interface | |
iface = gr.Interface( | |
fn=generate_code, | |
inputs=gr.Textbox(lines=5, label="Input Text"), | |
outputs=gr.Textbox(label="Generated Code"), | |
title="StarCoder Code Generator", | |
description="Enter a prompt to generate code using the StarCoder model.", | |
) | |
# Launch the Gradio interface | |
iface.launch() | |