import gradio as gr
from transformers import pipeline

# Load a completely free model (Falcon-1B)
pipe = pipeline(
    "text-generation",
    model="tiiuae/falcon-1b"  # ✅ Open, lightweight model that fits in free tier
)

# Function to generate AI responses
def generate_response(prompt):
    response = pipe(prompt, max_length=300, do_sample=True)
    return response[0]["generated_text"]

# Simple Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 NYT Connections Solver (Free AI Model)")

    input_text = gr.Textbox(label="Enter NYT Puzzle Words (Comma Separated)")
    output_text = gr.Textbox(label="AI's Categorization")

    def solve_puzzle(words):
        prompt = f"These words belong to four groups based on common themes: {words}. Identify the four groups."
        return generate_response(prompt)

    submit_btn = gr.Button("Solve Puzzle")
    submit_btn.click(solve_puzzle, inputs=input_text, outputs=output_text)

demo.launch()