Create app.py

app.py

@@ -0,0 +1,128 @@
+import gradio as gr
+import os
+import pymongo
+import spaces
+
+
+from huggingface_hub import login
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+def get_embedding(text: str) -> list[float]:
+    if not text.strip():
+        print("Attempted to get embedding for empty text.")
+        return []
+
+    embedding = embedding_model.encode(text)
+
+    return embedding.tolist()
+
+
+def get_mongo_client(mongo_uri):
+    """Establish connection to the MongoDB."""
+    try:
+        client = pymongo.MongoClient(mongo_uri)
+        print("Connection to MongoDB successful")
+        return client
+    except pymongo.errors.ConnectionFailure as e:
+        print(f"Connection failed: {e}")
+        return None
+
+
+def vector_search(user_query, collection):
+
+    # Generate embedding for the user query
+    query_embedding = get_embedding(user_query)
+
+    if query_embedding is None:
+        return "Invalid query or embedding generation failed."
+
+    # Define the vector search pipeline
+    pipeline = [
+        {
+            "$vectorSearch": {
+                "index": "vector_index",
+                "queryVector": query_embedding,
+                "path": "embedding",
+                "numCandidates": 150,  # Number of candidate matches to consider
+                "limit": 4,  # Return top 4 matches
+            }
+        },
+        {
+            "$project": {
+                "_id": 0,
+                "title": 1,
+                "ingredients": 1,
+                "directions": 1,
+                "score": {"$meta": "vectorSearchScore"},  # Include the search score
+            }
+        },
+    ]
+
+    # Execute the search
+    results = collection.aggregate(pipeline)
+    return list(results)
+
+
+def get_search_result(query, collection):
+
+    get_knowledge = vector_search(query, collection)
+
+    search_result = ""
+    for result in get_knowledge:
+        search_result += f"Recipe Name: {result.get('title', 'N/A')}, Ingredients: {result.get('ingredients', 'N/A')}, Directions: {result.get('directions', 'N/A')}\n"
+
+    return search_result, get_knowledge
+
+
+@spaces.GPU
+def process_response(message, history):
+  source_information, matches = get_search_result(message, collection)
+  recipe_dict = {}
+  for x in matches:
+    name = x.pop("title")
+    recipe_dict[name] = x
+
+  combined_information = f"Query: {message}\nContinue to answer the query by using the Search Results:\n{source_information}."
+  input_ids = tokenizer(combined_information, return_tensors="pt").to("cuda")
+  response = model.generate(**input_ids, max_new_tokens=500)
+  response_text = tokenizer.decode(response[0]).split("\n.\n")[-1].split("<eos>")[0].strip()
+
+  matched_recipe = ""
+  for title in recipe_dict.keys():
+    if title in response_text:
+      matched_recipe = title
+      break
+  if not matched_recipe:
+    matched_recipe = next(iter(recipe_dict))
+  recipe = recipe_dict[matched_recipe]
+
+  response_text += f"\n\nRecipe for **{matched_recipe}**:"
+  response_text += "\n### List of ingredients:\n- {0}".format("\n- ".join(recipe["ingredients"].split(", ")))
+  response_text += "\n### Directions:\n- {0}".format(".\n- ".join(recipe["directions"].split(". ")))
+
+  return response_text
+
+
+if __name__ == "__main__":
+    
+    # https://huggingface.co/thenlper/gte-large
+    embedding_model = SentenceTransformer("thenlper/gte-large")
+
+    mongo_uri = os.getenv("MONGO_URI")
+    if not mongo_uri:
+        raise ValueError("MONGO_URI not set in environment variables")
+
+    mongo_client = get_mongo_client(mongo_uri)
+    
+    # Ingest data into MongoDB
+    db = mongo_client["recipe"]
+    collection = db["recipe_collection"]
+
+    # login(token=os.getenv("HF_TOKEN"))
+
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
+    model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", device_map="auto")
+    
+    gr.ChatInterface(process_response).launch()