import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the model and tokenizer
model_path = "./trained_model"  # Current directory where model files are located
model = AutoModelForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# Define the function to generate a recipe
def generate_recipe(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs, 
        max_length=200, 
        temperature=0.8, 
        top_k=50, 
        top_p=0.95, 
        num_return_sequences=1
    )
    recipe = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return recipe

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_recipe,
    inputs=gr.Textbox(lines=2, placeholder="Enter a recipe prompt (e.g., 'Chocolate cake recipe')", label="Recipe Prompt"),
    outputs=gr.Textbox(label="Generated Recipe"),
    title="Recipe Generator",
    description="Enter a recipe idea or prompt, and let the AI generate a creative recipe for you!"
)

# Launch the app
interface.launch()