import os
import torch
import gradio as gr
import spaces  # REQUIRED FOR ZEROGPU
from transformers import AutoTokenizer, AutoModelForCausalLM  # Fixed: Added this missing import!

model_id = "mr-checker/GemmaExtract"
hf_token = os.getenv("HF_TOKEN")  

# Load Tokenizer & Model
tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16, 
    token=hf_token
).to("cuda") # ZeroGPU requires models to be pushed to cuda at the root level

system_prompt = (
    "You are an expert data analyst. Extract metadata from the YouTube title and description. "
    "Output MUST be a valid JSON object with the exact keys: 'market_signal', 'user_intent', and 'skills'."
)

# This tells HF to allocate a GPU when the button is clicked
@spaces.GPU  
def extract_metadata(title, description):
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Title: {title}\nDescription: {description}"}
    ]
    
    model_inputs = tokenizer.apply_chat_template(
        messages, 
        tokenize=True, 
        add_generation_prompt=True, 
        return_dict=True, 
        return_tensors="pt"
    ).to("cuda")

    outputs = model.generate(
        **model_inputs, 
        max_new_tokens=200,
        do_sample=False
    )

    input_len = model_inputs["input_ids"].shape[-1]
    generated_tokens = outputs[0][input_len:]
    return tokenizer.decode(generated_tokens, skip_special_tokens=True)

# Build Gradio UI
demo = gr.Interface(
    fn=extract_metadata,
    inputs=[
        gr.Textbox(label="YouTube Video Title", placeholder="Enter title..."),
        gr.Textbox(label="YouTube Video Description", placeholder="Enter description...", lines=4)
    ],
    outputs=gr.Code(label="Extracted JSON Output", language="json"),
    title="GemmaExtract - YouTube Metadata Extractor"
)

demo.launch()