ThoughtFocusAI/TF-CodeGenerator-Small

1import gradio as gr2import os3import torch4from transformers import RobertaTokenizer, T5ForConditionalGeneration5 6model_name = "ThoughtFocusAI/CodeGeneration-CodeT5-small"7device = "cuda" if torch.cuda.is_available() else "cpu"8model = T5ForConditionalGeneration.from_pretrained(9    model_name).to(device)10tokenizer = RobertaTokenizer.from_pretrained(11    model_name)12 13 14def generate_code(user_input):15    query = "Generate Python: " + user_input16    encoded_text = tokenizer(query, return_tensors='pt', padding='max_length',17                             truncation=True, max_length=512).input_ids.to(device)18 19    # inference20    generated_code = model.generate(encoded_text, max_length=512)21 22    # decode generated tokens23    decoded_code = tokenizer.decode(24        generated_code.numpy()[0], skip_special_tokens=True)25 26    return decoded_code27 28 29interface = gr.Interface(fn=generate_code,30                         inputs=gr.inputs.Textbox(31                             lines=3, label="Enter Text", placeholder="Ex-Add two numbers"),32                         outputs=gr.outputs.Textbox(label="Generated Code"))33interface.launch()34