Update app.py

app.py

@@ -18,7 +18,6 @@ data = dataset["train"]
 
 # Convert the string embeddings to numerical arrays and ensure they are 2D
 def convert_and_ensure_2d_embeddings(example):
-    # Clean the embedding string
     embedding_str = example['embedding']
     embedding_str = embedding_str.replace('\n', ' ').replace('...', '')
     embedding_list = list(map(float, embedding_str.strip("[]").split()))
@@ -44,11 +43,11 @@ data = data.map(flatten_embeddings)
 embeddings = np.array([example['embedding'] for example in data], dtype=np.float32)
 
 # Add FAISS index
-data.add_faiss_index_from_external_arrays("embedding", embeddings)
+data = data.add_faiss_index_from_external_arrays("embedding", embeddings)
 
 model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
 
-# use quantization to lower GPU usage
+# Use quantization to lower GPU usage
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
 )
@@ -71,16 +70,16 @@ You are given the extracted parts of a long document and a question. Provide a c
 If you don't know the answer, just say "I do not know." Don't make up an answer."""
 
 def search(query: str, k: int = 3):
-    """a function that embeds a new query and returns the most probable results"""
-    embedded_query = ST.encode(query)  # embed new query
-    scores, retrieved_examples = data.get_nearest_examples(  # retrieve results
-        "embedding", embedded_query,  # compare our new embedded query with the dataset embeddings
-        k=k  # get only top k results
+    """A function that embeds a new query and returns the most probable results."""
+    embedded_query = ST.encode(query)  # Embed new query
+    scores, retrieved_examples = data.get_nearest_examples(  # Retrieve results
+        "embedding", embedded_query,  # Compare our new embedded query with the dataset embeddings
+        k=k  # Get only top k results
     )
     return scores, retrieved_examples
 
 def format_prompt(prompt, retrieved_documents, k):
-    """using the retrieved documents we will prompt the model to generate our responses"""
+    """Using the retrieved documents we will prompt the model to generate our responses."""
     PROMPT = f"Question:{prompt}\nContext:"
     for idx in range(k):
         PROMPT += f"{retrieved_documents['text'][idx]}\n"
@@ -88,12 +87,12 @@ def format_prompt(prompt, retrieved_documents, k):
 
 @spaces.GPU(duration=150)
 def talk(prompt, history):
-    k = 1  # number of retrieved documents
+    k = 1  # Number of retrieved documents
     scores, retrieved_documents = search(prompt, k)
     formatted_prompt = format_prompt(prompt, retrieved_documents, k)
-    formatted_prompt = formatted_prompt[:2000]  # to avoid GPU OOM
+    formatted_prompt = formatted_prompt[:2000]  # To avoid GPU OOM
     messages = [{"role": "system", "content": SYS_PROMPT}, {"role": "user", "content": formatted_prompt}]
-    # tell the model to generate
+    # Tell the model to generate
     input_ids = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
@@ -131,12 +130,12 @@ def talk(prompt, history):
 TITLE = "# RAG"
 
 DESCRIPTION = """
-A rag pipeline with a chatbot feature
+A RAG pipeline with a chatbot feature
 Resources used to build this project :
-* embedding model : https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1
-* dataset : https://huggingface.co/datasets/not-lain/wikipedia
-* faiss docs : https://huggingface.co/docs/datasets/v2.18.0/en/package_reference/main_classes#datasets.Dataset.add_faiss_index 
-* chatbot : https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct
+* Embedding model : https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1
+* Dataset : https://huggingface.co/datasets/not-lain/wikipedia
+* FAISS docs : https://huggingface.co/docs/datasets/v2.18.0/en/package_reference/main_classes#datasets.Dataset.add_faiss_index 
+* Chatbot : https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct
 """
 
 demo = gr.ChatInterface(
@@ -150,7 +149,7 @@ demo = gr.ChatInterface(
         bubble_full_width=False,
     ),
     theme="Soft",
-    examples=[["what's anarchy ? "]],
+    examples=[["what's anarchy?"]],
     title=TITLE,
     description=DESCRIPTION,
 )