Update app.py

app.py

@@ -2,12 +2,10 @@ import streamlit as st
 import pandas as pd
 import chromadb
 from sentence_transformers import SentenceTransformer
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer
 from PIL import Image
 from io import BytesIO
 import requests
-from huggingface_hub import login
-
 
 # --- 1. Load Recipes Dataset ---
 @st.cache_data
@@ -29,7 +27,7 @@ recipes_df = load_recipes()
 # --- 2. Load SentenceTransformer Model ---
 @st.cache_resource
 def load_embedding_model():
-    return SentenceTransformer("all-mpnet-base-v2")
+    return SentenceTransformer("all-MiniLM-L6-v2")  # Smaller & optimized model
 
 embedding_model = load_embedding_model()
 
@@ -43,7 +41,7 @@ def get_sentence_transformer_embeddings(text):
 
 try:
     existing_data = collection.get()
-    existing_ids = set(existing_data["ids"]) if existing_data and "ids" in existing_data else set()
+    existing_ids = set(existing_data.get("ids", []))  # Use `.get()` for safety
 except Exception as e:
     st.error(f"⚠ ChromaDB Error: {e}")
     existing_ids = set()
@@ -54,78 +52,61 @@ for index, row in recipes_df.iterrows():
         continue
     embedding = get_sentence_transformer_embeddings(row["combined_text"])
     if embedding:
-        collection.add(embeddings=[embedding], documents=[row["combined_text"]], ids=[recipe_id])
+        collection.add(
+            embeddings=[embedding],
+            documents=[row["combined_text"]],
+            ids=[recipe_id],
+            metadatas=[{"title": row["title"], "ingredients": row["ingredients"], "instructions": row["instructions"], "img_src": row["img_src"]}]
+        )
 
 # --- 5. Retrieve Similar Recipes ---
 def retrieve_recipes(query, top_k=3):
     query_embedding = get_sentence_transformer_embeddings(query)
     results = collection.query(query_embeddings=[query_embedding], n_results=top_k)
-    
-    if results and "ids" in results and results["ids"] and results["ids"][0]:
+
+    if results and results.get("ids"):
         recipe_indices = [int(id) for id in results["ids"][0] if id.isdigit()]
         return recipes_df.iloc[recipe_indices] if recipe_indices else None
     return None
 
-
 hf_token = st.secrets["key"]
 if hf_token is None:
     raise ValueError("Hugging Face token is missing. Add it as a secret in your Space.")
-login(token=hf_token)   
-   
-# --- 6. Load Mistral-7B-Instruct ---
-@st.cache_resource
+login(token=hf_token)  
+
+# --- 6. Load LLM Model (Better Model for Generation) ---
 @st.cache_resource
-def load_mistral_model():
-    model_name = "mistralai/Mistral-7B-Instruct-v0.3"
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
-    model = AutoModelForCausalLM.from_pretrained(
-    model_name, use_auth_token=True, trust_remote_code=True, device_map="auto")
-    return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150)
-
-mistral_model = load_mistral_model()
-
-
-# --- 7. Answer Question Using Mistral ---
-def answer_question(query, context=""):
-    greetings = ["hi", "hello", "hey", "greetings", "how are you", "what's up"]
-    query_cleaned = query.lower().strip()
-
-    # Handle greetings
-    if query_cleaned in greetings:
-        return "Hello! I'm here to assist with recipes and food-related questions. 🍽️ What would you like to know?"
-
-    # Retrieve relevant recipe
-    related_recipes = retrieve_recipes(query, top_k=1)
-    if related_recipes is None or related_recipes.empty:
-        return "I specialize in recipes! 🍽️ Feel free to ask me about ingredients, cooking methods, or meal ideas. 😊"
-
-    # If found, use its instructions as context
-    context = related_recipes.iloc[0]['instructions']
-    prompt = f"Context: {context}\n\nQuestion: {query}\nAnswer:"
-    
-    response = mistral_model(prompt)
-    if isinstance(response, list) and response:
-        return response[0].get("generated_text", "I'm not sure, but I can help with recipes! 😊").strip()
-    
-    return "I'm not sure, but I can help with recipes! 😊"
-
-# --- 8. Classify Query Type ---
+def load_llm_model():
+    return pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.3")
+
+llm_model = load_llm_model()
+
+# --- 7. Load Classifier ---
 @st.cache_resource
 def load_classifier():
-    return pipeline("zero-shot-classification", model="facebook/bart-large-mnli", use_auth_token=True)
+    return pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
 classifier = load_classifier()
 
-def classify_query(query):
-    recipe_keywords = ["make", "cook", "bake", "recipe", "prepare"]
-    if any(keyword in query.lower() for keyword in recipe_keywords):
-        return "Recipe Search"
-    
-    labels = ["Q&A", "Recipe Search"]
+# --- 8. Query Classification ---
+def is_food_related(query):
+    labels = ["Food & Cooking", "General Knowledge", "Science", "Technology"]
     result = classifier(query, candidate_labels=labels, multi_label=False)
-    return result.get("labels", ["Q&A"])[0]
+    return result["labels"][0] == "Food & Cooking"
+
+def generate_recipe(query):
+    related_recipes = retrieve_recipes(query, top_k=2)
 
-# --- 9. Display Image ---
+    if not related_recipes or related_recipes.empty:
+        return "I couldn't find a matching recipe, but let me create one for you!"
+
+    base_text = "\n".join([f"- {r['title']}: {r['ingredients']}" for _, r in related_recipes.iterrows()])
+    full_prompt = f"Create a unique recipe using these ingredients: {query}.\n\nReference recipes:\n{base_text}"
+
+    response = llm_model(full_prompt, max_length=200, num_return_sequences=1)
+    return response[0]["generated_text"]
+
+# --- 9. Display Image Function ---
 def display_image(image_url, recipe_name):
     try:
         if not isinstance(image_url, str) or not image_url.startswith("http"):
@@ -140,7 +121,7 @@ def display_image(image_url, recipe_name):
         st.image(placeholder_url, caption=recipe_name, use_container_width=True)
 
 # --- 10. Streamlit UI ---
-st.title("🍽️ AI Recipe & Q&A Assistant (Powered by Mistral-7B)")
+st.title("🍽️ AI Recipe & Q&A Assistant")
 
 user_query = st.text_input("Enter your question or recipe search query:", "", key="main_query_input")
 
@@ -149,31 +130,22 @@ if "retrieved_recipes" not in st.session_state:
 
 if st.button("Ask AI"):
     if user_query:
-        # Handle greetings separately
-        greeting_response = answer_question(user_query)
-        if greeting_response.startswith("Hello!"):
-            st.subheader("🤖 AI Answer:")
-            st.write(greeting_response)
-        else:
-            # Classify query
-            intent = classify_query(user_query)
-
-            if intent == "Q&A":
-                st.subheader("🤖 AI Answer:")
-                response = answer_question(user_query)
-                st.write(response)
-
-            elif intent == "Recipe Search":
-                retrieved_recipes = retrieve_recipes(user_query)
-                if retrieved_recipes is not None and not retrieved_recipes.empty:
-                    st.session_state["retrieved_recipes"] = retrieved_recipes
-                    st.subheader("🍴 Found Recipes:")
-                    for index, recipe in retrieved_recipes.iterrows():
-                        st.markdown(f"### {recipe['title']}")
-                        st.write(f"**Ingredients:** {recipe['ingredients']}")
-                        st.write(f"**Instructions:** {recipe['instructions']}")
-                        display_image(recipe.get('img_src', ''), recipe['title'])
-                else:
-                    st.warning("⚠️ No relevant recipes found.")
+        if is_food_related(user_query):  # Check if it's food-related
+            st.subheader("🤖 AI Response:")
+            response = generate_recipe(user_query)
+            st.write(response)
+
+            retrieved_recipes = retrieve_recipes(user_query)
+            if retrieved_recipes is not None and not retrieved_recipes.empty:
+                st.session_state["retrieved_recipes"] = retrieved_recipes
+                st.subheader("🍴 Found Recipes:")
+                for _, recipe in retrieved_recipes.iterrows():
+                    st.markdown(f"### {recipe['title']}")
+                    st.write(f"**Ingredients:** {recipe['ingredients']}")
+                    st.write(f"**Instructions:** {recipe['instructions']}")
+                    display_image(recipe.get('img_src', ''), recipe['title'])
             else:
-                st.warning("❌ Unable to classify the query.")
+                st.warning("⚠️ No relevant recipes found.")
+        else:
+            st.subheader("🤖 AI Answer:")
+            st.write("I specialize in food-related topics! 🍽️ Please ask me about recipes, cooking methods, or ingredients.")