import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "Manvtith/FoodExtract-gemma-3-270m-fine-tune-v1"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto"
)

def extract_food(text):
    prompt = f"Extract all food and drink items from this sentence:\n{text}\nAnswer:"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=128,
            temperature=0.2,
            do_sample=False
        )

    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result.split("Answer:")[-1].strip()

demo = gr.Interface(
    fn=extract_food,
    inputs=gr.Textbox(lines=3, placeholder="Type a sentence with food..."),
    outputs="text",
    title="Food & Drink Extractor (Gemma SLM)",
    description="Fine-tuned Gemma-3-270M to extract food and beverage items from text."
)

demo.launch()