Create app.py

app.py

@@ -0,0 +1,110 @@
+import faiss
+import numpy as np
+from FlagEmbedding import FlagModel
+from flask import Flask, request, jsonify
+from datasets import load_dataset
+import gradio as gr
+import os
+import time
+from functools import lru_cache
+
+# Initialize components
+app = Flask(__name__)
+model = None
+index = None
+corpus = None
+
+def initialize_components():
+    global model, index, corpus
+    
+    # Load model with safety checks
+    if model is None:
+        model = FlagModel(
+            "BAAI/bge-large-en-v1.5",
+            query_instruction_for_retrieval="Represent this sentence for searching relevant passages:",
+            use_fp16=True
+        )
+    
+    # Load corpus from Hugging Face dataset
+    if corpus is None:
+        dataset = load_dataset("awinml/medrag_corpus_sampled", split='train')
+        corpus = [f"{row['id']}\t{row['contents']}" for row in dataset]
+    
+    # Create FAISS index in memory
+    if index is None:
+        embeddings = model.encode([doc.split('\t', 1)[1] for doc in corpus])
+        dimension = embeddings.shape[1]
+        index = faiss.IndexFlatIP(dimension)
+        index.add(embeddings.astype('float32'))
+
+@app.route("/retrieve", methods=["POST"])
+def retrieve():
+    start_time = time.time()
+    
+    # Validate request
+    data = request.json
+    if not data or "queries" not in data:
+        return jsonify({"error": "Missing 'queries' parameter"}), 400
+    
+    # Initialize components if needed
+    initialize_components()
+    
+    # Process queries
+    queries = data["queries"]
+    topk = data.get("topk", 3)
+    return_scores = data.get("return_scores", False)
+    
+    # Batch processing
+    query_embeddings = model.encode_queries(queries)
+    scores, indices = index.search(query_embeddings.astype('float32'), topk)
+    
+    # Format results
+    results = []
+    for i, query in enumerate(queries):
+        query_results = []
+        for j in range(topk):
+            doc_idx = indices[i][j]
+            doc = corpus[doc_idx]
+            doc_id, content = doc.split('\t', 1)
+            result = {
+                "document": {
+                    "id": doc_id,
+                    "contents": content
+                },
+                "score": float(scores[i][j])
+            }
+            query_results.append(result)
+        results.append(query_results)
+    
+    return jsonify({
+        "result": results,
+        "time": f"{time.time() - start_time:.2f}s"
+    })
+
+# Gradio UI for testing
+def gradio_interface(query, topk):
+    response = requests.post(
+        "http://localhost:7860/retrieve",
+        json={"queries": [query], "topk": topk}
+    )
+    return response.json()["result"][0]
+
+# Start server
+if __name__ == "__main__":
+    # First-time initialization
+    initialize_components()
+    
+    # Create Gradio interface
+    iface = gr.Interface(
+        fn=gradio_interface,
+        inputs=[
+            gr.Textbox(label="Medical Query", placeholder="Enter your medical question..."),
+            gr.Slider(1, 10, value=3, label="Top Results")
+        ],
+        outputs=gr.JSON(label="Retrieval Results"),
+        title="Medical Retrieval System",
+        description="Search across medical literature using AI-powered semantic search"
+    )
+    
+    # Run both Flask and Gradio
+    iface.launch(server_name="0.0.0.0", server_port=7860, share=True)