"""
Quick inference script for Gemma-3-270M Football Extractor
"""

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import json

def extract_football_data(post_text: str, model_path: str = "."):
    """
    Extract structured data from a football news post.
    
    Args:
        post_text: The football news post text
        model_path: Path to the model directory
    
    Returns:
        dict: Extracted structured data
    """
    # Load model and tokenizer
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16,
        device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    # Prepare messages
    messages = [
        {
            "role": "system",
            "content": "You are a silent, expert data extraction API. Respond ONLY with JSON in markdown code block."
        },
        {
            "role": "user",
            "content": f"Please extract the structured data from the following post:\n\n{post_text}"
        }
    ]
    
    # Generate
    inputs = tokenizer.apply_chat_template(
        messages,
        return_tensors="pt",
        add_generation_prompt=True
    ).to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            max_new_tokens=1024,
            temperature=0.1,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id
        )
    
    # Decode and parse
    result = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
    
    # Extract JSON
    if "```json" in result:
        json_str = result.split("```json")[1].split("```")[0].strip()
    else:
        json_str = result.strip()
    
    return json.loads(json_str)

# Example usage
if __name__ == "__main__":
    post = "🚨 BREAKING: Manchester United agree £80m deal for Jude Bellingham!"
    result = extract_football_data(post)
    print(json.dumps(result, indent=2))