CodeTuneStudio / app.py
whackthejacker's picture
Create app.py
b81292f verified
raw
history blame
1.38 kB
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()