import gradio as gr
from transformers import T5ForConditionalGeneration, T5Tokenizer
import os

# Global variables
model = None
tokenizer = None

def load_model():
    global model, tokenizer
    if model is None:
        print("Loading model...")
        tokenizer = T5Tokenizer.from_pretrained("younsmirzai/stackadvisor-extractor")
        model = T5ForConditionalGeneration.from_pretrained("younsmirzai/stackadvisor-extractor")
        print("✅ Model loaded!")
    return model, tokenizer

def extract_order(ocr_text):
    if not ocr_text or not ocr_text.strip():
        return "Error: Empty input"
    
    try:
        m, t = load_model()
        inputs = t(ocr_text, return_tensors="pt", max_length=256, truncation=True)
        outputs = m.generate(**inputs, max_length=128)
        result = t.decode(outputs[0], skip_special_tokens=True)
        return result
    except Exception as e:
        return f"Error: {str(e)}"

# Simple interface
demo = gr.Interface(
    fn=extract_order,
    inputs=gr.Textbox(label="OCR Text", lines=3),
    outputs=gr.Textbox(label="Result"),
    title="StackAdvisor",
)

demo.launch()