khairo14's picture
Update app.py
943b15b verified
Raw
History Blame Contribute Delete
457 Bytes
import gradio as gr
from transformers import pipeline
# You can change the model to any light-weight code model
pipe = pipeline("text-generation", model="Salesforce/codegen-350M-mono")
def generate(prompt):
result = pipe(prompt, max_new_tokens=200)
return result[0]['generated_text']
gr.Interface(
fn=generate,
inputs=gr.Textbox(lines=4, placeholder="Enter code prompt here..."),
outputs="text",
title="Code Helper AI"
).launch()