Update app.py

app.py

@@ -8,7 +8,7 @@ from haystack.schema import Document
 
 from typing import Optional, List
 
-from huggingface_hub import get_inference_endpoint
+# from huggingface_hub import get_inference_endpoint
 from datasets import load_dataset
 from time import perf_counter
 import gradio as gr
@@ -16,18 +16,27 @@ import numpy as np
 import requests
 import os
 
+TOP_K = 2
+BATCH_SIZE = 16
 
-RETRIEVER_URL = os.getenv("RETRIEVER_URL")
-RANKER_URL = os.getenv("RANKER_URL")
 HF_TOKEN = os.getenv("HF_TOKEN")
+RANKER_URL = os.getenv("RANKER_URL")
+RETRIEVER_URL = os.getenv("RETRIEVER_URL")
 
+# 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
+# )
 
-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
-)
+# if RETRIEVER_IE.status != "running":
+#     RETRIEVER_IE.resume()
+#     RETRIEVER_IE.wait()
+
+# if RANKER_IE.status != "running":
+#     RANKER_IE.resume()
+#     RANKER_IE.wait()
 
 
 def post(url, payload):
@@ -137,10 +146,6 @@ class Ranker(BaseRanker):
         return [[Document.from_dict(d) for d in docs] for docs in response]
 
 
-TOP_K = 2
-BATCH_SIZE = 16
-
-
 if (
     os.path.exists("/data/faiss_document_store.db")
     and os.path.exists("/data/faiss_index.json")
@@ -152,21 +157,27 @@ if (
     )
     document_store.save(index_path="/data/faiss_index")
 else:
-    try:
-        os.remove("/data/faiss_index")
-        os.remove("/data/faiss_index.json")
-        os.remove("/data/faiss_document_store.db")
-    except FileNotFoundError:
-        pass
+    for file in [
+        "/data/faiss_document_store.db",
+        "/data/faiss_index.json",
+        "/data/faiss_index",
+    ]:
+        try:
+            os.remove(file)
+        except FileNotFoundError:
+            pass
 
     document_store = FAISSDocumentStore(
         sql_url="sqlite:////data/faiss_document_store.db",
         return_embedding=True,
         embedding_dim=384,
     )
-    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.write_documents(
+        load_dataset("bilgeyucel/seven-wonders", split="train")
+    )
+    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")
 
@@ -178,27 +189,9 @@ pipe.add_node(component=ranker, name="Ranker", inputs=["Retriever"])
 
 
 def run(query: str) -> dict:
-    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>
-    """
+    output = f"""<h2>Top {TOP_K} Documents</h2>"""
 
     for i, doc in enumerate(pipe_output["documents"]):
         output += f"""
@@ -221,23 +214,24 @@ examples = [
     "How did Colossus of Rhodes collapse?",
 ]
 
-
 input_text = gr.components.Textbox(
-    label="Query",
-    placeholder="Enter a query",
-    value=examples[0],
-    lines=3,
+    label="Query", placeholder="Enter a query", value=examples[0], lines=1
 )
-output_html = gr.components.HTML(label="Results")
+output_html = gr.components.HTML(label="Documents")
 
 gr.Interface(
     fn=run,
     inputs=input_text,
     outputs=output_html,
-    title="End-to-End Retrieval & Ranking",
     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).",
+    cache_examples=False,
+    allow_flagging="never",
+    title="End-to-End Retrieval & Ranking with Hugging Face Inference Endpoints and Spaces",
+    description="""## A [haystack](https://haystack.deepset.ai/) pipeline with the following components
+- <strong>Retriever</strong>: [Quantized FastRAG Retriever](https://huggingface.co/optimum-intel/fastrag-retriever) deployed on [Inference Endpoints](https://huggingface.co/docs/inference-endpoints/index) + Intel Sapphire Rapids CPU.
+- <strong>Ranker</strong>: [Quantized FastRAG Retriever](https://huggingface.co/optimum-intel/fastrag-ranker) deployed on [Inference Endpoints](https://huggingface.co/docs/inference-endpoints/index) + Intel Sapphire Rapids CPU.
+- <strong>Document Store</strong>: A [FAISS document store](https://github.com/facebookresearch/faiss/tree/main) containing the [`seven-wonders` dataset](https://huggingface.co/datasets/bilgeyucel/seven-wonders), created on this Space's [persistent storage](https://huggingface.co/docs/hub/en/spaces-storage).
+
+This Space is based on the optimizations demonstrated in the blog [CPU Optimized Embeddings with 🤗 Optimum Intel and fastRAG](https://huggingface.co/blog/intel-fast-embedding)
+""",
 ).launch()