Spaces:
Runtime error
Runtime error
from haystack.document_stores.faiss import FAISSDocumentStore | |
from haystack.nodes.retriever import EmbeddingRetriever | |
from haystack.nodes.ranker import BaseRanker | |
from haystack.pipelines import Pipeline | |
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( | |
url, | |
json=payload, | |
headers={"Authorization": f"Bearer {HF_TOKEN}"}, | |
) | |
return response.json() | |
def method_timer(method): | |
def timed(self, *args, **kw): | |
start_time = perf_counter() | |
result = method(self, *args, **kw) | |
end_time = perf_counter() | |
print( | |
f"{self.__class__.__name__}.{method.__name__} took {end_time - start_time} seconds" | |
) | |
return result | |
return timed | |
class Retriever(EmbeddingRetriever): | |
def __init__( | |
self, | |
document_store: Optional[BaseDocumentStore] = None, | |
top_k: int = 10, | |
batch_size: int = 32, | |
scale_score: bool = True, | |
): | |
self.document_store = document_store | |
self.top_k = top_k | |
self.batch_size = batch_size | |
self.scale_score = scale_score | |
def embed_queries(self, queries: List[str]) -> np.ndarray: | |
payload = {"queries": queries, "inputs": ""} | |
response = post(RETRIEVER_URL, payload) | |
if "error" in response: | |
raise gr.Error(response["error"]) | |
arrays = np.array(response) | |
return arrays | |
def embed_documents(self, documents: List[Document]) -> np.ndarray: | |
documents = [d.to_dict() for d in documents] | |
for doc in documents: | |
doc["embedding"] = None | |
payload = {"documents": documents, "inputs": ""} | |
response = post(RETRIEVER_URL, payload) | |
if "error" in response: | |
raise gr.Error(response["error"]) | |
arrays = np.array(response) | |
return arrays | |
class Ranker(BaseRanker): | |
def predict( | |
self, query: str, documents: List[Document], top_k: Optional[int] = None | |
) -> List[Document]: | |
documents = [d.to_dict() for d in documents] | |
for doc in documents: | |
doc["embedding"] = None | |
payload = {"query": query, "documents": documents, "top_k": top_k, "inputs": ""} | |
response = post(RANKER_URL, payload) | |
if "error" in response: | |
raise gr.Error(response["error"]) | |
return [Document.from_dict(d) for d in response] | |
def predict_batch( | |
self, | |
queries: List[str], | |
documents: List[List[Document]], | |
batch_size: Optional[int] = None, | |
top_k: Optional[int] = None, | |
) -> List[List[Document]]: | |
documents = [[d.to_dict() for d in docs] for docs in documents] | |
for docs in documents: | |
for doc in docs: | |
doc["embedding"] = None | |
payload = { | |
"queries": queries, | |
"documents": documents, | |
"batch_size": batch_size, | |
"top_k": top_k, | |
"inputs": "", | |
} | |
response = post(RANKER_URL, payload) | |
if "error" in response: | |
raise gr.Error(response["error"]) | |
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") | |
and os.path.exists("/data/faiss_index") | |
): | |
document_store = FAISSDocumentStore.load("/data/faiss_index") | |
retriever = Retriever( | |
document_store=document_store, top_k=TOP_K, batch_size=BATCH_SIZE | |
) | |
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 | |
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.update_embeddings(retriever=retriever) | |
document_store.save(index_path="/data/faiss_index") | |
ranker = Ranker() | |
pipe = Pipeline() | |
pipe.add_node(component=retriever, name="Retriever", inputs=["Query"]) | |
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> | |
""" | |
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?", | |
] | |
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=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).", | |
).launch() | |