services/classify.py

from __future__ import annotations

from functools import lru_cache
from typing import Literal

import torch
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline

from .config import settings
from .text_utils import normalize_text

RecipeType = Literal["baked", "cooked"]

BAKE_KEYWORDS = [
    "bake", "baking", "oven", "preheat", "flour", "dough", "batter",
    "cake", "cookie", "muffin", "bread", "pastry", "brownie", "tart",
    "pie", "scone", "loaf", "whisk", "fold in", "sift", "knead",
    "leavening", "baking soda", "baking powder", "yeast",
]
COOK_KEYWORDS = [
    "saute", "sauté", "fry", "boil", "simmer", "stir", "grill",
    "roast", "steam", "poach", "braise", "sear", "stove", "skillet",
    "pan", "wok", "sauce", "soup", "stew", "marinate",
]


@lru_cache(maxsize=1)
def get_qa_pipeline():
    tokenizer = AutoTokenizer.from_pretrained(settings.qa_model_name)
    model = AutoModelForQuestionAnswering.from_pretrained(settings.qa_model_name)
    device = 0 if torch.cuda.is_available() else -1
    return pipeline(
        "question-answering",
        model=model,
        tokenizer=tokenizer,
        device=device,
    )


def classify_recipe(recipe_text: str) -> RecipeType:
    text = normalize_text(recipe_text)

    bake_score = sum(1 for kw in BAKE_KEYWORDS if kw in text)
    cook_score = sum(1 for kw in COOK_KEYWORDS if kw in text)

    answer = ""
    try:
        qa = get_qa_pipeline()
        result = qa(question="Is this recipe for baking or cooking?", context=recipe_text)
        answer = normalize_text(str(result.get("answer", "")))
    except Exception:
        pass

    if any(sig in answer for sig in ("bak", "oven", "pastry", "dough")):
        return "baked"
    if any(sig in answer for sig in ("cook", "fry", "boil", "saut", "grill", "stir")):
        return "cooked"

    if bake_score > cook_score:
        return "baked"
    if cook_score > bake_score:
        return "cooked"

    return "cooked"