tayyab / model.py
mabdullahsibghatullah123's picture
Upload 17 files
5ae7c8f verified
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
import pickle
import os
# 1. Dataset of common foods and their typical shelf life (in days)
# This acts as our "Knowledge Base"
data = [
# Dairy
("Milk", 7), ("Yogurt", 14), ("Cheese", 30), ("Butter", 60), ("Cream", 10),
# Vegetables
("Spinach", 5), ("Lettuce", 5), ("Tomato", 7), ("Carrot", 21), ("Potato", 30), ("Onion", 30),
# Fruits
("Apple", 21), ("Banana", 5), ("Orange", 14), ("Grapes", 7), ("Strawberry", 4),
# Meat
("Chicken", 2), ("Beef", 3), ("Pork", 3), ("Fish", 2), ("Ham", 5),
# Grains
("Bread", 5), ("Rice", 365), ("Pasta", 365), ("Cereal", 180),
# Misc
("Eggs", 21), ("Juice", 10), ("Sauce", 90), ("Canned Beans", 700)
]
MODEL_PATH = "expiry_model.pkl"
def train_model():
print("Training Expiry Prediction Model...")
df = pd.DataFrame(data, columns=["item", "days"])
# Simple NLP pipeline: Text -> Vector -> Classifier
model = make_pipeline(TfidfVectorizer(), MultinomialNB())
model.fit(df["item"], df["days"])
with open(MODEL_PATH, "wb") as f:
pickle.dump(model, f)
print("Model saved to", MODEL_PATH)
return model
def load_model():
if os.path.exists(MODEL_PATH):
with open(MODEL_PATH, "rb") as f:
return pickle.load(f)
else:
return train_model()
def predict_days(item_name):
model = load_model()
try:
# Predict
predicted_days = model.predict([item_name])[0]
return int(predicted_days)
except Exception as e:
print(f"Prediction Error: {e}")
return 7 # Default fallback