Minor details for RAG

app.py

@@ -185,13 +185,13 @@ def generate_cells(dataset_id, cells, notebook_type="eda"):
         cells, wildcards, replacements, has_numeric_columns, has_categoric_columns
     )
     generated_text = ""
-    # Show only the first 40 lines, would like to have a scroll in gr.Code https://github.com/gradio-app/gradio/issues/9192
+    # Show only the first 30 lines, would like to have a scroll in gr.Code https://github.com/gradio-app/gradio/issues/9192
     for cell in cells:
         if cell["cell_type"] == "markdown":
             continue
         generated_text += cell["source"] + "\n\n"
         yield generated_text, ""
-        if generated_text.count("\n") > 38:
+        if generated_text.count("\n") > 30:
             generated_text += (
                 f"## See more lines available in the generated notebook 🤗 ......"
             )

utils/notebook_utils.py

@@ -291,15 +291,16 @@ rag_cells = [
         "cell_type": "code",
         "source": """
 # Install and import necessary libraries.
-!pip install pandas sentence-transformers faiss-cpu transformers torch
+!pip install pandas sentence-transformers faiss-cpu transformers torch huggingface_hub
 """,
     },
     {
         "cell_type": "code",
         "source": """
-import pandas as pd
 from sentence_transformers import SentenceTransformer
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from huggingface_hub import InferenceClient
+import pandas as pd
 import faiss
 import torch
 """,
@@ -364,15 +365,10 @@ index.add(vectors)
         "cell_type": "code",
         "source": """
 # Specify the text you want to search for in the list
-text_to_search = text_list[0]
-print(f"Text to search: {text_to_search}")
-""",
-    },
-    {
-        "cell_type": "code",
-        "source": """
+query = "How to prepare a cake?"
+
 # Generate the embedding for the search query
-query_embedding = model.encode([text_to_search])
+query_embedding = model.encode([query])
 """,
     },
     {
@@ -385,7 +381,10 @@ D, I = index.search(query_embedding, k=10)
 print(f"Similar documents: {[text_list[i] for i in I[0]]}")
 """,
     },
-    {"cell_type": "markdown", "source": "## 4. Load pipeline and perform inference"},
+    {
+        "cell_type": "markdown",
+        "source": "## 4. Load pipeline and perform inference locally",
+    },
     {
         "cell_type": "code",
         "source": """
@@ -404,19 +403,43 @@ generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device
         "cell_type": "code",
         "source": """
 # Create a prompt with two parts: 'system' for instructions based on a 'context' from the retrieved documents, and 'user' for the query
-query = "How to prepare a cake?"
 selected_elements = [text_list[i] for i in I[0].tolist()]
 context = ','.join(selected_elements)
-prompt = f"system: Answer user's question based on '{context}'. user: {query}"
+messages = [
+    {
+        "role": "system",
+        "content": f"You are an intelligent assistant tasked with providing accurate and concise answers based on the following context. Use the information retrieved to construct your response. Context: {context}",
+    },
+    {"role": "user", "content": query},
+]
 """,
     },
     {
         "cell_type": "code",
         "source": """
 # Send the prompt to the pipeline and show the answer
-output = generator(prompt)
-print("Generated Summary:")
-print(output[0]['generated_text'])
+output = generator(messages)
+print("Generated result:")
+print(output[0]['generated_text'][-1]['content']) # Print the assistant's response content
+""",
+    },
+    {
+        "cell_type": "markdown",
+        "source": "## 5. Alternatively call the inference client",
+    },
+    {
+        "cell_type": "code",
+        "source": """
+# Adjust model name as needed
+checkpoint = "meta-llama/Meta-Llama-3-8B-Instruct"
+
+# Change here your Hugging Face API token
+token = "hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" 
+
+inference_client = InferenceClient(checkpoint, token=token)
+output = inference_client.chat_completion(messages=messages, stream=False)
+print("Generated result:")
+print(output.choices[0].message.content)
 """,
     },
 ]