import gradio as gr
from transformers import pipeline

# Load the model
model = pipeline("text-generation", model="Salesforce/codegen-350M-multi", max_new_tokens=512)

# Define the function that gets called from JS
def generate(prompt):
    result = model(f"# HTML/CSS/JS request: {prompt}\n")[0]["generated_text"]
    return result

# Serve API endpoint at /generate
gr.Interface(fn=generate, inputs="text", outputs="text").launch(
    inline=False,
    server_name="0.0.0.0",
    server_port=7860
)