Spaces:
Running
Running
import gradio as gr | |
import torch | |
from transformers import CodeT5ForConditionalGeneration, CodeT5Tokenizer | |
# Load pre-trained CodeT5 model and tokenizer | |
model = CodeT5ForConditionalGeneration.from_pretrained("Salesforce/code-t5-small") | |
tokenizer = CodeT5Tokenizer.from_pretrained("Salesforce/code-t5-small") | |
def generate_code(prompt, code_file): | |
# Read uploaded code file | |
if code_file: | |
code_text = code_file.read().decode("utf-8") | |
else: | |
code_text = "" | |
# Tokenize input prompt | |
input_ids = tokenizer.encode(prompt, return_tensors="pt") | |
# Generate code using CodeT5 model | |
output = model.generate(input_ids=input_ids, max_length=256) | |
generated_code = tokenizer.decode(output[0], skip_special_tokens=True) | |
# Return generated code and code preview | |
return generated_code, f"```python\n{generated_code}\n```" | |
# Create Gradio interface | |
iface = gr.Interface( | |
fn=generate_code, | |
inputs=[ | |
________gr.Textbox(label="Input_Prompt",_placeholder="Enter_a_prompt"), | |
________gr.Upload(label="Upload_Code_File",_file_types=["py"]) | |
], | |
outputs=[ | |
________gr.Textbox(label="Generated_Code"), | |
________gr.Code(label="Code_Preview",_language="python") | |
____], | |
title="Code Generation with CodeT5", | |
description="Generate Python code based on input prompt and uploaded code file." | |
) | |
# Launch Gradio interface | |
iface.launch() | |