app.py

"""
World Cup Storyteller — Hugging Face Spaces app.py
NLP Homework 4 — ARI 525

Three approaches, equal treatment:
- Zero-Shot : prompt only → LLM knowledge
- Few-Shot  : prompt + 2 style examples → LLM knowledge
- RAG       : prompt → retrieve from WC dataset → grounded generation
"""

import os
import time
import warnings
warnings.filterwarnings("ignore")

import pandas as pd
import numpy as np
import gradio as gr
from groq import Groq
from sentence_transformers import SentenceTransformer
import faiss

# -----------------------------------------------
# CONFIG
# -----------------------------------------------

GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "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")
MODEL = "llama-3.1-8b-instant"
client = Groq(api_key=GROQ_API_KEY)

# -----------------------------------------------
# LOAD DATA (used only by RAG)
# -----------------------------------------------

try:
    cups    = pd.read_csv("WorldCups.csv")
    matches = pd.read_csv("WorldCupMatches.csv")
    print("✅ Dataset loaded.")
except FileNotFoundError as e:
    raise RuntimeError(
        "CSV files not found. Upload WorldCups.csv and WorldCupMatches.csv "
        "to the Space root folder."
    ) from e


# -----------------------------------------------
# DATA CLEANING & CONTEXT BUILDERS
# -----------------------------------------------

def clean_data(cups, matches):
    matches_clean = matches.dropna(
        subset=["Home Team Name", "Away Team Name",
                "Home Team Goals", "Away Team Goals"]
    ).copy()
    matches_clean["Home Team Name"] = matches_clean["Home Team Name"].str.strip()
    matches_clean["Away Team Name"]  = matches_clean["Away Team Name"].str.strip()
    matches_clean["Home Team Goals"] = matches_clean["Home Team Goals"].astype(int)
    matches_clean["Away Team Goals"] = matches_clean["Away Team Goals"].astype(int)
    cups_clean = cups.dropna(subset=["Year", "Winner"]).copy()
    cups_clean["Year"] = cups_clean["Year"].astype(int)
    return cups_clean, matches_clean


def build_team_context(team, year, cups, matches):
    cup_info = cups[cups["Year"] == year]
    if cup_info.empty:
        return None
    cup = cup_info.iloc[0]
    team_matches = matches[
        (matches["Year"] == year) &
        ((matches["Home Team Name"] == team) | (matches["Away Team Name"] == team))
    ].sort_values("Stage")
    if team_matches.empty:
        return None

    match_lines = []
    for _, row in team_matches.iterrows():
        home, away = row["Home Team Name"], row["Away Team Name"]
        hg, ag = int(row["Home Team Goals"]), int(row["Away Team Goals"])
        if home == team:
            result = "WIN" if hg > ag else ("DRAW" if hg == ag else "LOSS")
            line = f"[{row['Stage']}] {team} vs {away}: {hg}-{ag} ({result})"
        else:
            result = "WIN" if ag > hg else ("DRAW" if hg == ag else "LOSS")
            line = f"[{row['Stage']}] {team} vs {home}: {ag}-{hg} ({result})"
        match_lines.append(line)

    final_note = ""
    if cup.get("Winner") == team:
        final_note = f"{team} WON the {year} World Cup! 🏆"
    elif cup.get("Runners-Up") == team:
        final_note = f"{team} were runners-up in {year}."
    elif cup.get("Third") == team:
        final_note = f"{team} finished third in {year}."

    return (
        f"{team}'s Journey — World Cup {year}\n"
        f"Host: {cup.get('Country', 'Unknown')}\n"
        f"{final_note}\n\nMatch-by-match results:\n" +
        "\n".join(match_lines)
    )


def build_edition_context(year, cups, matches):
    cup_info = cups[cups["Year"] == year]
    if cup_info.empty:
        return None
    cup = cup_info.iloc[0]
    edition_matches = matches[matches["Year"] == year].sort_values("Stage")
    match_lines = [
        f"[{row['Stage']}] {row['Home Team Name']} {int(row['Home Team Goals'])} "
        f"- {int(row['Away Team Goals'])} {row['Away Team Name']}"
        for _, row in edition_matches.iterrows()
    ]
    return (
        f"World Cup {year} — Host: {cup.get('Country', 'Unknown')}\n"
        f"Winner: {cup.get('Winner', 'Unknown')}\n"
        f"Runners-up: {cup.get('Runners-Up', 'Unknown')}\n"
        f"Third Place: {cup.get('Third', 'Unknown')}\n"
        f"Goals Scored: {cup.get('GoalsScored', 'Unknown')}\n"
        f"Attendance: {cup.get('Attendance', 'Unknown')}\n\nAll Matches:\n" +
        "\n".join(match_lines)
    )


cups_clean, matches_clean = clean_data(cups, matches)


# -----------------------------------------------
# RAG KNOWLEDGE BASE
# -----------------------------------------------

def build_knowledge_base(cups, matches):
    chunks, metadata = [], []
    for _, cup_row in cups.iterrows():
        year = int(cup_row["Year"])
        year_matches = matches[matches["Year"] == year]
        teams = set(
            year_matches["Home Team Name"].tolist() +
            year_matches["Away Team Name"].tolist()
        )
        for team in teams:
            ctx = build_team_context(team, year, cups, matches)
            if ctx:
                chunks.append(ctx)
                metadata.append({"team": team, "year": year})
        ctx = build_edition_context(year, cups, matches)
        if ctx:
            chunks.append(ctx)
            metadata.append({"team": "ALL", "year": year})
    return chunks, metadata


print("⏳ Loading embedding model...")
embedder = SentenceTransformer("all-MiniLM-L6-v2")

print("⏳ Building & embedding knowledge base...")
kb_chunks, kb_metadata = build_knowledge_base(cups_clean, matches_clean)
kb_embeddings = embedder.encode(kb_chunks, show_progress_bar=False, convert_to_numpy=True)

dim = kb_embeddings.shape[1]
faiss_index = faiss.IndexFlatL2(dim)
faiss_index.add(kb_embeddings)
print(f"✅ RAG index ready — {faiss_index.ntotal} vectors")


def retrieve_context(query, top_k=3):
    query_vec = embedder.encode([query], convert_to_numpy=True)
    distances, indices = faiss_index.search(query_vec, top_k)
    retrieved = []
    for idx in indices[0]:
        if idx < len(kb_chunks):
            meta = kb_metadata[idx]
            retrieved.append(f"[Retrieved: {meta['team']} — {meta['year']}]\n{kb_chunks[idx]}")
    return "\n\n---\n\n".join(retrieved)


# -----------------------------------------------
# PROMPTS & FEW-SHOT EXAMPLES
# -----------------------------------------------

BASE_SYSTEM_PROMPT = """
You are a passionate, knowledgeable sports journalist specializing in FIFA World Cup history.
Generate vivid, engaging, narrative-driven stories about World Cup tournaments and team journeys.

Guidelines:
- Write in flowing narrative prose — no bullet points
- Build tension and highlight dramatic moments
- Be factually accurate based on your knowledge
- Adapt tone to the user's request (documentary, dramatic, casual, etc.)
- Keep stories between 250-400 words unless asked otherwise
""".strip()

RAG_SYSTEM_PROMPT = """
You are a passionate sports journalist specializing in FIFA World Cup history.
You will be given retrieved historical match data. Use it to write an accurate,
vivid, narrative-driven story. Only use facts present in the retrieved data.

Guidelines:
- Write in flowing narrative prose — no bullet points
- Build tension and highlight dramatic moments
- Stick strictly to the facts in the retrieved data — do not invent scores
- Adapt tone to the user's request
- Keep stories between 250-400 words unless asked otherwise
""".strip()

FEW_SHOT_EXAMPLES = [
    {
        "prompt": "Tell the story of France's 1998 World Cup victory.",
        "story": (
            "It was the summer that France found its destiny on home soil. Les Bleus entered "
            "the 1998 World Cup as hosts carrying immense pressure, but from the very first "
            "whistle they played with a quiet, relentless authority.\n\n"
            "The group stage was a statement — South Africa swept aside 3-0, Saudi Arabia "
            "dismantled 4-0, Denmark edged out 2-1. By the knockout rounds, France carried "
            "the weight of a nation's expectations into every match.\n\n"
            "Paraguay made them suffer — a lone golden goal was all that separated the sides. "
            "Italy pushed them to penalties, a nerve-shredding duel France survived with "
            "ice-cold nerves. Then came Croatia. France trailed before turning the game on "
            "its head to win 2-1.\n\n"
            "The final against defending champions Brazil became a coronation. Three goals, "
            "zero reply. The Stade de France erupted, and a generation of French children "
            "discovered what it felt like to be champions of the world."
        )
    },
    {
        "prompt": "Tell the story of West Germany's 1954 World Cup win.",
        "story": (
            "They called it the Miracle of Bern, and for good reason. No one believed West "
            "Germany could win the 1954 World Cup — least of all after Hungary handed them "
            "an 8-3 humiliation in the group stage.\n\n"
            "But West Germany, crafty and resilient, had rested key players for that match "
            "and quietly plotted their path forward. They dispatched Turkey twice, squeezed "
            "past Yugoslavia, then demolished Austria 6-1 in a dazzling semi-final.\n\n"
            "The final was a rematch nobody expected. Hungary — the Mighty Magyars, unbeaten "
            "for four years — led 2-0 within eight minutes. The world assumed it was over.\n\n"
            "It was not. West Germany clawed back to 2-2, and then six minutes from the end "
            "Helmut Rahn struck. 3-2. A country still rebuilding from the rubble of war had "
            "become world champions."
        )
    }
]


# -----------------------------------------------
# GENERATION FUNCTIONS
# -----------------------------------------------

def generate_zeroshot(user_prompt):
    start = time.time()
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": BASE_SYSTEM_PROMPT},
            {"role": "user",   "content": user_prompt}
        ],
        temperature=0.8,
        max_tokens=600
    )
    elapsed = round(time.time() - start, 2)
    return response.choices[0].message.content.strip(), elapsed


def generate_fewshot(user_prompt):
    messages = [{"role": "system", "content": BASE_SYSTEM_PROMPT}]
    for ex in FEW_SHOT_EXAMPLES:
        messages.append({"role": "user",      "content": ex["prompt"]})
        messages.append({"role": "assistant", "content": ex["story"]})
    messages.append({"role": "user", "content": user_prompt})
    start = time.time()
    response = client.chat.completions.create(
        model=MODEL, messages=messages, temperature=0.8, max_tokens=600
    )
    elapsed = round(time.time() - start, 2)
    return response.choices[0].message.content.strip(), elapsed


def generate_rag(user_prompt, top_k=3):
    retrieved = retrieve_context(user_prompt, top_k=top_k)
    full_message = (
        f"Retrieved World Cup data:\n---\n{retrieved}\n---\n\n"
        f"User request: {user_prompt}"
    )
    start = time.time()
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": RAG_SYSTEM_PROMPT},
            {"role": "user",   "content": full_message}
        ],
        temperature=0.8,
        max_tokens=600
    )
    elapsed = round(time.time() - start, 2)
    return response.choices[0].message.content.strip(), elapsed


# -----------------------------------------------
# GRADIO UI
# -----------------------------------------------

def generate_story_ui(user_prompt, approach):
    if not user_prompt.strip():
        return "⚠️ Please enter a prompt.", ""
    try:
        if approach == "Zero-Shot":
            story, elapsed = generate_zeroshot(user_prompt)
            info = f"⚡ Zero-Shot (LLM knowledge only) | ⏱️ {elapsed}s | 📝 {len(story.split())} words"
        elif approach == "Few-Shot":
            story, elapsed = generate_fewshot(user_prompt)
            info = f"📖 Few-Shot (LLM knowledge + 2 style examples) | ⏱️ {elapsed}s | 📝 {len(story.split())} words"
        else:
            story, elapsed = generate_rag(user_prompt)
            info = f"🔍 RAG (retrieved from WC dataset) | ⏱️ {elapsed}s | 📝 {len(story.split())} words"
        return story, info
    except Exception as e:
        return f"❌ Error: {str(e)}", ""


EXAMPLES = [
    "Tell me the story of Brazil's 2002 World Cup campaign like a sports documentary.",
    "What happened in the 1966 World Cup? Who won and how?",
    "Describe Argentina's 1986 World Cup triumph — focus on the drama.",
    "Tell Italy's 2006 World Cup story.",
    "What was the most dramatic World Cup final ever?",
]

with gr.Blocks(
    title="⚽ World Cup Storyteller",
    theme=gr.themes.Base(),
    css="""
        #header { text-align: center; padding: 1.5em 0 0.5em; }
        #header h1 { font-size: 2.2em; }
        #header p  { color: #777; font-size: 1.05em; }
        #story-box textarea { font-size: 1.05em; line-height: 1.8; }
        .note { font-size: 0.85em; color: #777; margin-top: 0.4em; line-height: 1.6; }
    """
) as demo:

    with gr.Column(elem_id="header"):
        gr.Markdown("# ⚽ World Cup Storyteller")
        gr.Markdown(
            "Ask anything about World Cup history — get a vivid narrative story. "
            "Compare how different NLP approaches handle the same prompt."
        )

    with gr.Row():

        with gr.Column(scale=1, min_width=260):
            gr.Markdown("### ⚙️ NLP Approach")
            approach = gr.Radio(
                choices=["Zero-Shot", "Few-Shot", "RAG"],
                value="Few-Shot",
                label=""
            )
            gr.Markdown(
                "**Zero-Shot** — Prompt only. No examples, no data. "
                "The LLM relies entirely on its own knowledge.\n\n"
                "**Few-Shot** — 2 hand-crafted story examples guide the writing style. "
                "Still uses LLM knowledge, no dataset.\n\n"
                "**RAG** — Retrieves real match data from the WC dataset (1930–2014) "
                "before generating. Most factually grounded.",
                elem_classes="note"
            )

        with gr.Column(scale=2):
            gr.Markdown("### ✍️ Your Prompt")
            user_prompt = gr.Textbox(
                placeholder="e.g. Tell me the story of Brazil's 2002 World Cup campaign...",
                label="Ask anything about World Cup history",
                lines=3
            )
            gr.Examples(
                examples=[[p] for p in EXAMPLES],
                inputs=user_prompt,
                label="Try an example"
            )
            btn = gr.Button("🎙️ Generate Story", variant="primary", size="lg")

            gr.Markdown("### 📖 Story")
            story_out = gr.Textbox(
                label="", lines=16, interactive=False, elem_id="story-box"
            )
            info_out = gr.Markdown("")

    btn.click(
        fn=generate_story_ui,
        inputs=[user_prompt, approach],
        outputs=[story_out, info_out]
    )

    gr.Markdown(
        "---\n"
        "*Data: FIFA World Cup dataset (1930–2014) · "
        "LLM: Llama 3 8B via Groq · "
        "Embeddings: all-MiniLM-L6-v2 · "
        "Built for NLP HW4 — ARI 525*"
    )

demo.launch()