Update app.py

app.py

@@ -7,17 +7,28 @@ from haystack.document_stores.base import BaseDocumentStore
 from haystack.schema import Document
 
 from typing import Optional, List
+
+from huggingface_hub import get_inference_endpoint
+from datasets import load_dataset
 from time import perf_counter
 import gradio as gr
 import numpy as np
 import requests
 import os
 
+
 RETRIEVER_URL = os.getenv("RETRIEVER_URL")
 RANKER_URL = os.getenv("RANKER_URL")
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 
+RETRIEVER_IE = get_inference_endpoint(
+    "fastrag-retriever", namespace="optimum-intel", token=HF_TOKEN
+)
+RANKER_IE = get_inference_endpoint(
+    "fastrag-ranker", namespace="optimum-intel", token=HF_TOKEN
+)
+
 
 def post(url, payload):
     response = requests.post(
@@ -128,16 +139,7 @@ class Ranker(BaseRanker):
 
 TOP_K = 2
 BATCH_SIZE = 16
-EXAMPLES = [
-    "There is a blue house on Oxford Street.",
-    "Paris is the capital of France.",
-    "The Eiffel Tower is in Paris.",
-    "The Louvre is in Paris.",
-    "London is the capital of England.",
-    "Cairo is the capital of Egypt.",
-    "The pyramids are in Egypt.",
-    "The Sphinx is in Egypt.",
-]
+
 
 if (
     os.path.exists("/data/faiss_document_store.db")
@@ -148,7 +150,6 @@ if (
     retriever = Retriever(
         document_store=document_store, top_k=TOP_K, batch_size=BATCH_SIZE
     )
-    document_store.update_embeddings(retriever=retriever)
     document_store.save(index_path="/data/faiss_index")
 else:
     try:
@@ -163,12 +164,9 @@ else:
         return_embedding=True,
         embedding_dim=384,
     )
-    document_store.write_documents(
-        [Document(content=d, id=i) for i, d in enumerate(EXAMPLES)]
-    )
-    retriever = Retriever(
-        document_store=document_store, top_k=TOP_K, batch_size=BATCH_SIZE
-    )
+    DATASET = load_dataset("bilgeyucel/seven-wonders", split="train")
+    document_store.write_documents(DATASET)
+    retriever = Retriever(document_store=document_store, top_k=TOP_K, batch_size=BATCH_SIZE)
     document_store.update_embeddings(retriever=retriever)
     document_store.save(index_path="/data/faiss_index")
 
@@ -180,20 +178,66 @@ pipe.add_node(component=ranker, name="Ranker", inputs=["Retriever"])
 
 
 def run(query: str) -> dict:
-    output = pipe.run(query=query)
-    closest_documents = [d.content for d in output["documents"]]
-    return f"Closest ({TOP_K}) document(s): {closest_documents}"
+    if RETRIEVER_IE.status != "running":
+        RETRIEVER_IE.resume()
+        raise gr.Error(
+            "Retriever Inference Endpoint is not running. "
+            "Sent a request to resume it. Please try again in a few minutes."
+        )
+
+    if RANKER_IE.status != "running":
+        RANKER_IE.resume()
+        raise gr.Error(
+            "Ranker Inference Endpoint is not running. "
+            "Sent a request to resume it. Please try again in a few minutes."
+        )
+
+    pipe_output = pipe.run(query=query)
+
+    output = f"""
+    <h2>Query</h2>
+    <p>{query}</p>
+    <h2>Top {TOP_K} Documents</h2>
+    """
+
+    for i, doc in enumerate(pipe_output["documents"]):
+        output += f"""
+        <h3>Document {i + 1}</h3>
+        <p><strong>ID:</strong> {doc.id}</p>
+        <p><strong>Score:</strong> {doc.score}</p>
+        <p><strong>Content:</strong> {doc.content}</p>
+        """
+
+    return output
+
+
+examples = [
+    "Where is Gardens of Babylon?",
+    "Why did people build Great Pyramid of Giza?",
+    "What does Rhodes Statue look like?",
+    "Why did people visit the Temple of Artemis?",
+    "What is the importance of Colossus of Rhodes?",
+    "What happened to the Tomb of Mausolus?",
+    "How did Colossus of Rhodes collapse?",
+]
 
 
-run("What is the capital of France?")
-print("Warmed up successfully!")
+input_text = gr.components.Textbox(
+    label="Query",
+    placeholder="Enter a query",
+    value=examples[0],
+    lines=3,
+)
+output_html = gr.components.HTML(label="Results")
 
 gr.Interface(
     fn=run,
-    inputs="text",
-    outputs="text",
+    inputs=input_text,
+    outputs=output_html,
     title="End-to-End Retrieval & Ranking",
-    examples=["What is the capital of France?"],
-    description="A pipeline for retrieving and ranking documents "
-    "from a memory persistent FAISS document store, using Inference Endpoints.",
+    examples=examples,
+    description="A [haystack](https://haystack.deepset.ai/) pipeline for retrieving and ranking "
+    "documents from the [seven-wonders dataset](bilgeyucel/seven-wonders) based on a query, "
+    "using a FAISS database as a document store (kept in the space's persistent storage) "
+    "and two [Inference Endpoints for the Retriever and Ranker](https://huggingface.co/collections/optimum-intel/fast-rag-inference-endpoints-6641c6cbb98ddf3fe49c7728).",
 ).launch()