Spaces:
Build error
Build error
| 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) | |
| # inference | |
| generated_code = model.generate(encoded_text, max_length=512) | |
| # decode generated tokens | |
| 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() | |