Update app.py

app.py

@@ -5,8 +5,20 @@ handling the Gradio/Streamlit UI and orchestrating the RAG pipeline
 to answer queries based on the game's mechanics and manual.
 """
 import gradio as gr
-from helpers import create_prompt, find_similar_documents, load_vector_store
-from model import CHAT_MODEL_ID, generate_answer, load_chat_model
+from helpers import (
+    create_prompt, 
+    find_similar_documents, 
+    load_vector_store, 
+    CHROMA_PATH,
+    COLLECTION_NAME,
+)
+from model import (
+    CHAT_MODEL_ID, 
+    EMBEDDING_MODEL_ID,
+    generate_answer, 
+    load_chat_model, 
+    load_embedding_model,
+)
 import pytest
 import sys
 
@@ -20,7 +32,11 @@ else:
 # Initialiations and globals
 print(f"Loading {CHAT_MODEL_ID} model...")
 chat_model, tokenizer = load_chat_model()
-print("Model loaded!")
+print("Chat model loaded!")
+
+print(f"Loading {EMBEDDING_MODEL_ID} model...")
+embedding_model = load_embedding_model()
+print("Embedding model loaded!"
 
 MANUAL_PATH = "emerald_manual.txt"
 print(f"Loading vector store from {MANUAL_PATH}...")
@@ -29,7 +45,8 @@ print("Vector store loaded!")
 
 def respond(message, history):
     print("Enter app.respond...")
-    context = find_similar_documents(vector_store, message)
+    context = find_similar_documents(
+        CHROMA_PATH, COLLECTION_NAME, embedding_model, message, k=5)
     prompt = create_prompt(message, history, context)
     response = generate_answer(chat_model, tokenizer, prompt)
     if not response: