File size: 1,978 Bytes
df8f88e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | 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"
|