Create app.py

app.py

@@ -0,0 +1,78 @@
+import gradio as gr
+from huggingface_hub import get_token, InferenceClient
+from sentence_transformers import CrossEncoder
+
+inference_client = InferenceClient(api_key=get_token())
+reranker = CrossEncoder("sentence-transformers/all-MiniLM-L12-v2")
+
+
+def query_and_rerank_documents(query: str, k_retrieved: int = 10):
+    documents = similarity_search(query, k_retrieved)
+    documents = documents.drop_duplicates("chunk")
+    documents["rank"] = reranker.predict([[query, hit] for hit in documents["chunk"]])
+    reranked_documents = documents.sort_values(by="rank", ascending=False)
+    return reranked_documents
+
+
+def generate_response_api(query: str):
+    messages = [
+        {
+            "role": "system",
+            "content": "You will receive a query and context. Only return the answer based on the context without mentioning the context.",
+        },
+        {"role": "user", "content": query},
+    ]
+    completion = inference_client.chat.completions.create(
+        model="HuggingFaceTB/SmolLM2-360M-Instruct", messages=messages, max_tokens=2000
+    )
+
+    return completion.choices[0].message
+
+
+def rag_pipeline(query: str, k_retrieved: int = 10, k_reranked: int = 5):
+    documents = query_and_rerank_documents(query, k_retrieved=k_retrieved)
+    query_with_context = (
+        f"Context: {documents['chunk'].to_list()[:k_reranked]}\n\nQuery: {query}"
+    )
+    return generate_response_api(query_with_context), documents
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("""# RAG Hub Datasets 
+                
+                Part of [smol blueprint](https://github.com/davidberenstein1957/smol-blueprint) - a smol blueprint for AI development, focusing on practical examples of RAG, information extraction, analysis and fine-tuning in the age of LLMs.""")
+
+    with gr.Row():
+        query_input = gr.Textbox(
+            label="Query", placeholder="Enter your question here...", lines=3
+        )
+
+    with gr.Row():
+        with gr.Column():
+            retrieve_slider = gr.Slider(
+                minimum=1,
+                maximum=20,
+                value=10,
+                label="Number of documents to retrieve",
+            )
+        with gr.Column():
+            rerank_slider = gr.Slider(
+                minimum=1,
+                maximum=10,
+                value=5,
+                label="Number of documents to use after reranking",
+            )
+
+    submit_btn = gr.Button("Submit")
+    response_output = gr.Textbox(label="Response", lines=10)
+    documents_output = gr.Dataframe(
+        label="Documents", headers=["chunk", "url", "distance", "rank"], wrap=True
+    )
+
+    submit_btn.click(
+        fn=rag_interface,
+        inputs=[query_input, retrieve_slider, rerank_slider],
+        outputs=[response_output, documents_output],
+    )
+
+demo.launch()