|
|
import gradio as gr |
|
|
import os |
|
|
import torch |
|
|
from transformers import RobertaTokenizer, T5ForConditionalGeneration |
|
|
|
|
|
model_name = "ThoughtFocusAI/CodeGeneration-CodeT5-small" |
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
model = T5ForConditionalGeneration.from_pretrained( |
|
|
model_name).to(device) |
|
|
tokenizer = RobertaTokenizer.from_pretrained( |
|
|
model_name) |
|
|
|
|
|
|
|
|
def generate_code(user_input): |
|
|
query = "Generate Python: " + user_input |
|
|
encoded_text = tokenizer(query, return_tensors='pt', padding='max_length', |
|
|
truncation=True, max_length=512).input_ids.to(device) |
|
|
|
|
|
|
|
|
generated_code = model.generate(encoded_text, max_length=512) |
|
|
|
|
|
|
|
|
decoded_code = tokenizer.decode( |
|
|
generated_code.numpy()[0], skip_special_tokens=True) |
|
|
|
|
|
return decoded_code |
|
|
|
|
|
|
|
|
interface = gr.Interface(fn=generate_code, |
|
|
inputs=gr.inputs.Textbox( |
|
|
lines=3, label="Enter Text", placeholder="Ex-Add two numbers"), |
|
|
outputs=gr.outputs.Textbox(label="Generated Code")) |
|
|
interface.launch() |
|
|
|
