Upload inference.py with huggingface_hub

inference.py

@@ -1,48 +1,77 @@
 """
-Simple inference script for the weight predictor model.
-Usage:
-    python inference.py "cheeseburger with fries" menu_item
-    python inference.py "chocolate bar 100g" grocery
-    python inference.py "wireless mouse" non_food
+Usage example for the hybrid weight predictor.
+
+Requirements:
+  pip install huggingface_hub joblib scikit-learn
+
+This uses a hybrid approach:
+  1. Explicit weight extraction from text (e.g., "500g", "2 liter")
+  2. Knowledge base lookup for known FMCG brands and fast food items
+  3. ML model fallback for unknown items
 """
-import sys
-import joblib
 from huggingface_hub import hf_hub_download
+import joblib
+import sys
+import os
 
-# Import predictor classes so pickling works
-import weight_predictor
+# Make sure the weight_predictor module is importable
+sys.path.insert(0, os.path.dirname(__file__))
 
+# Load the hybrid predictor (pure Python, no ML model needed)
+from hybrid_weight_predictor import HybridWeightPredictor, build_hybrid_predictor
 
-def load_model():
-    print("Loading model from HF Hub...")
-    model_path = hf_hub_download(
-        repo_id="ZZandro/weight-predictor",
-        filename="unified_predictor.pkl",
-        repo_type="model"
-    )
-    return joblib.load(model_path)
+def predict_weight(text, item_type=None):
+    """
+    Predict item weight in grams.
 
+    Args:
+        text: Item description. Should start with [MENU_ITEM], [GROCERY], or [NON_FOOD]
+        item_type: "menu_item", "grocery", or "non_food" (auto-detected from text if None)
 
-def predict(text, item_type="grocery"):
-    predictor = load_model()
-    weight = predictor.predict(text, item_type=item_type)
-    return weight
+    Returns:
+        Predicted weight in grams (float)
+    """
+    predictor = build_hybrid_predictor()
+    return predictor.predict(text, item_type)
 
 
 if __name__ == "__main__":
-    if len(sys.argv) < 2:
-        print("Usage: python inference.py '<item description>' <item_type>")
-        print("  item_type: menu_item | grocery | non_food")
-        sys.exit(1)
-
-    description = sys.argv[1]
-    item_type = sys.argv[2] if len(sys.argv) > 2 else "grocery"
-
-    # Auto-prepend tag if not present
-    if not description.startswith("["):
-        description = f"[{item_type.upper()}] {description}"
-
-    weight = predict(description, item_type)
-    print(f"\nItem: {sys.argv[1]}")
-    print(f"Type: {item_type}")
-    print(f"Predicted weight: {weight:.1f} grams")
+    # Example usage
+    examples = [
+        ("[GROCERY] coca cola can 330ml", "grocery"),
+        ("[GROCERY] coca cola", "grocery"),  # uses KB default
+        ("[GROCERY] pepsi 1 liter bottle", "grocery"),
+        ("[GROCERY] kelloggs corn flakes 500g", "grocery"),
+        ("[GROCERY] oreo cookies 154g", "grocery"),
+        ("[GROCERY] heinz ketchup 570ml", "grocery"),
+        ("[GROCERY] mars bar 51g", "grocery"),
+        ("[GROCERY] snickers 2 pack 96g", "grocery"),
+        ("[GROCERY] red bull 4 pack", "grocery"),
+        ("[GROCERY] tide laundry detergent 1.5kg", "grocery"),
+        ("[MENU_ITEM] large pizza", "menu_item"),
+        ("[MENU_ITEM] cheeseburger", "menu_item"),
+        ("[MENU_ITEM] double cheeseburger", "menu_item"),
+        ("[MENU_ITEM] big mac", "menu_item"),
+        ("[MENU_ITEM] french fries", "menu_item"),
+        ("[MENU_ITEM] large fries", "menu_item"),
+        ("[MENU_ITEM] chicken nuggets", "menu_item"),
+        ("[MENU_ITEM] burrito", "menu_item"),
+        ("[MENU_ITEM] caesar salad", "menu_item"),
+        ("[MENU_ITEM] caesar salad large", "menu_item"),
+        ("[MENU_ITEM] pho", "menu_item"),
+        ("[MENU_ITEM] ramen", "menu_item"),
+        ("[MENU_ITEM] sushi platter", "menu_item"),
+        ("[MENU_ITEM] medium pizza", "menu_item"),
+        ("[MENU_ITEM] combo meal", "menu_item"),
+        ("[MENU_ITEM] milkshake", "menu_item"),
+        ("[MENU_ITEM] iced coffee", "menu_item"),
+        ("[MENU_ITEM] family meal", "menu_item"),
+        ("[MENU_ITEM] sliders", "menu_item"),
+        ("[NON_FOOD] laptop computer", "non_food"),
+        ("[NON_FOOD] water bottle", "non_food"),
+    ]
+
+    print("=== Weight Predictions ===\n")
+    for text, item_type in examples:
+        weight = predict_weight(text, item_type)
+        print(f"  {text:55s} -> {weight:8.1f}g")