import gradio as gr
from transformers import pipeline
import re

# Load model
generator = pipeline(
    "text-generation",
    model="tiiuae/falcon-7b-instruct",
    use_auth_token="HF_TOKEN2"  # Replace with your real token
)

def parse_flashcards(text):
    # Extract Q/A pairs using regex
    qa_pairs = re.findall(r"Q[:\-]?\s*(.*?)\s*A[:\-]?\s*(.*?)(?=Q[:\-]|$)", text, re.DOTALL)
    flashcards = []
    for i, (q, a) in enumerate(qa_pairs, start=1):
        flashcards.append(f"Q{i}: {q.strip()}\nA{i}: {a.strip()}")
    return "\n\n".join(flashcards) if flashcards else text

def generate_flashcards(topic):
    prompt = f"Create 5 unique flashcards about '{topic}'. Format:\nQ: <question>\nA: <answer>\n"
    response = generator(
        prompt,
        max_new_tokens=300,
        temperature=0.7,
        top_p=0.9
    )
    raw_text = response[0]["generated_text"]
    return parse_flashcards(raw_text)

with gr.Blocks() as demo:
    gr.Markdown("## 🃏 AI Flashcard Generator\nEnter a topic to generate flashcards instantly.")
    with gr.Row():
        topic = gr.Textbox(label="Enter Topic", placeholder="e.g. Python basics")
        btn = gr.Button("Generate Flashcards")
    output = gr.Textbox(label="Flashcards", lines=15)
    btn.click(generate_flashcards, inputs=topic, outputs=output)

demo.launch()