Update app.py

app.py

@@ -1,204 +1,177 @@
-import os
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModel
 import numpy as np
-from typing import List, Tuple
+from typing import List, Dict
 import pandas as pd
-from huggingface_hub import login
+import os
+from pinecone import Pinecone
 
-login(os.getenv("HF_TOKEN"))
-# Model configuration
-MODEL_NAME = "ekacare/parrotlet-e"
 
-class ParrotletEmbedder:
-    def __init__(self, model_name: str):
-        """Initialize the Parrotlet-E model and tokenizer."""
+
+if PINECONE_API_KEY is None:
+    raise ValueError("Please set PINECONE_API_KEY as an environment variable.")
+
+
+# =========================
+# Retriever Class
+# =========================
+class ParrotletRetriever:
+    def __init__(self, model_name: str, pinecone_api_key: str, index_name: str):
+        """Initialize model and Pinecone client."""
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Loading model on {self.device}...")
-        
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+        print(f"🚀 Loading model on {self.device}...")
+
+        # Load tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
+        self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True, token=HF_TOKEN)
         self.model.to(self.device)
         self.model.eval()
-        
+
+        # Pinecone initialization (unchanged)
+        self.pinecone_namespace = os.environ.get("NAMESPACE")
+        self.pinecone_client = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))
+        self.pinecone_index = self.pinecone_client.Index(host=os.environ.get("PINECONE_HOST"))
+
+        print(f"Connected to Pinecone index: {index_name}")
+        if self.pinecone_namespace:
+            print(f"🔹 Using namespace: {self.pinecone_namespace}")
+
+    # --------------------------
+    # Mean Pooling
+    # --------------------------
     def mean_pooling(self, model_output, attention_mask):
-        """Perform mean pooling on model output."""
+        """Mean pooling for sentence embeddings."""
         token_embeddings = model_output[0]
         input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-    
-    def encode(self, texts: List[str], batch_size: int = 32) -> np.ndarray:
-        """Encode texts into embeddings."""
-        all_embeddings = []
-        
+        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
+            input_mask_expanded.sum(1), min=1e-9
+        )
+
+    # --------------------------
+    # Text Encoder
+    # --------------------------
+    def encode(self, texts: List[str]) -> np.ndarray:
+        """Encode text into normalized embeddings."""
         with torch.no_grad():
-            for i in range(0, len(texts), batch_size):
-                batch_texts = texts[i:i + batch_size]
-                
-                encoded_input = self.tokenizer(
-                    batch_texts,
-                    padding=True,
-                    truncation=True,
-                    max_length=512,
-                    return_tensors='pt'
-                ).to(self.device)
-                
-                model_output = self.model(**encoded_input)
-                embeddings = self.mean_pooling(model_output, encoded_input['attention_mask'])
-                
-                # Normalize embeddings
-                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
-                all_embeddings.append(embeddings.cpu().numpy())
-        
-        return np.vstack(all_embeddings)
-    
-    def compute_similarity(self, text1: str, text2: str) -> float:
-        """Compute cosine similarity between two texts."""
-        embeddings = self.encode([text1, text2])
-        similarity = np.dot(embeddings[0], embeddings[1])
-        return float(similarity)
-
-# Initialize model
-embedder = ParrotletEmbedder(MODEL_NAME)
-
-def compute_pairwise_similarity(query: str, documents: str) -> Tuple[pd.DataFrame, str]:
-    """
-    Compute similarity between query and multiple documents.
-    
-    Args:
-        query: The query text
-        documents: Documents separated by newlines
-    
-    Returns:
-        DataFrame with documents and similarity scores, and status message
-    """
+            encoded_input = self.tokenizer(
+                texts, padding=True, truncation=True, max_length=512, return_tensors="pt"
+            ).to(self.device)
+
+            model_output = self.model(**encoded_input)
+            embeddings = self.mean_pooling(model_output, encoded_input["attention_mask"])
+            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+            return embeddings.cpu().numpy()
+
+    # --------------------------
+    # Pinecone Search
+    # --------------------------
+    def search(self, query: str, top_k: int = 5) -> List[Dict]:
+        """Search Pinecone index."""
+        query_vector = self.encode([query])[0]
+
+        results = self.pinecone_index.query(
+            namespace=self.pinecone_namespace,
+            vector=query_vector.tolist(),
+            top_k=top_k,
+            include_metadata=True,
+            include_values=False,
+        )
+
+        docs = []
+        for i, match in enumerate(results["matches"]):
+            text = match["metadata"].get("text", "[No text metadata]")
+            docs.append({
+                "Rank": i + 1,
+                "Score": f"{match['score']:.4f}",
+                "Document": text
+                "ID": match["concept_id"],
+            })
+
+        return docs
+
+
+
+MODEL_NAME = "ekacare/parrotlet-e"
+retriever = ParrotletRetriever(MODEL_NAME, PINECONE_API_KEY, PINECONE_INDEX_NAME)
+
+
+def retrieve_documents(query: str, top_k: int = 5):
+    """Perform retrieval and return results."""
     if not query.strip():
-        return pd.DataFrame(), "⚠️ Please enter a query text."
-    
-    if not documents.strip():
-        return pd.DataFrame(), "⚠️ Please enter at least one document."
-    
-    # Split documents by newlines and filter empty lines
-    doc_list = [doc.strip() for doc in documents.split('\n') if doc.strip()]
-    
-    if len(doc_list) == 0:
-        return pd.DataFrame(), "⚠️ No valid documents found."
-    
+        return pd.DataFrame(), "⚠️ Please enter a valid query."
+
     try:
-        # Encode query and documents
-        query_embedding = embedder.encode([query])
-        doc_embeddings = embedder.encode(doc_list)
-        
-        # Compute similarities
-        similarities = np.dot(doc_embeddings, query_embedding.T).flatten()
-        
-        # Create results dataframe
-        results = pd.DataFrame({
-            'Rank': range(1, len(doc_list) + 1),
-            'Document': doc_list,
-            'Similarity Score': similarities
-        })
-        
-        # Sort by similarity
-        results = results.sort_values('Similarity Score', ascending=False).reset_index(drop=True)
-        results['Rank'] = range(1, len(results) + 1)
-        
-        status = f"✅ Successfully computed similarities for {len(doc_list)} documents."
-        return results, status
-        
+        results = retriever.search(query, top_k)
+        if not results:
+            return pd.DataFrame(), "⚠️ No results found."
+
+        df = pd.DataFrame(results)
+        status = f"✅ Retrieved top {len(results)} documents."
+        return df, status
+
     except Exception as e:
         return pd.DataFrame(), f"❌ Error: {str(e)}"
 
-def compute_single_similarity(text1: str, text2: str) -> Tuple[str, str]:
-    """
-    Compute similarity between two texts.
-    
-    Args:
-        text1: First text
-        text2: Second text
-    
-    Returns:
-        Similarity score and status message
-    """
-    if not text1.strip() or not text2.strip():
-        return "", "⚠️ Please enter both texts."
-    
-    try:
-        similarity = embedder.compute_similarity(text1, text2)
-        score_display = f"### Similarity Score: {similarity:.4f}"
-        status = "✅ Similarity computed successfully."
-        return score_display, status
-        
-    except Exception as e:
-        return "", f"❌ Error: {str(e)}"
-
-# Create Gradio interface
-with gr.Blocks(title="Parrotlet-e: Indic Medical Embedding Model", theme=gr.themes.Soft()) as demo:
-    
-    with gr.Tab("Query-Document Matching"):
-        gr.Markdown("""
-        ### 📄 Semantic Search""")
-        
-        with gr.Row():
-            with gr.Column():
-                query_input = gr.Textbox(
-                    label="term1",
-                    placeholder="",
-                    lines=1
-                )
-                documents_input = gr.Textbox(
-                    label="term2",
-                    placeholder="",
-                    lines=1
-                )
-                search_btn = gr.Button("🔍 Search", variant="primary")
-            
-            with gr.Column():
-                search_output = gr.Dataframe(
-                    label="Results",
-                    headers=["Rank", "Document", "Similarity Score"],
-                    datatype=["number", "str", "number"],
-                    wrap=True
-                )
-                search_status = gr.Textbox(label="Status", interactive=False)
-        
-        search_btn.click(
-            fn=compute_pairwise_similarity,
-            inputs=[query_input, documents_input],
-            outputs=[search_output, search_status]
-        )
-    
-    with gr.Tab("Pairwise Similarity"):
-        gr.Markdown("""
-        ### 🔗 Compare Two Texts
-        Compute semantic similarity between any two medical texts (score ranges from -1 to 1).
-        """)
-        
-        with gr.Row():
-            with gr.Column():
-                text1_input = gr.Textbox(
-                    label="Text 1",
-                    placeholder="Enter first text...\nExample: हृदय रोग के लक्षण",
-                    lines=5
-                )
-                text2_input = gr.Textbox(
-                    label="Text 2",
-                    placeholder="Enter second text...\nExample: छाती में दर्द और सांस फूलना",
-                    lines=5
-                )
-                similarity_btn = gr.Button("⚡ Calculate Similarity", variant="primary")
-            
-            with gr.Column():
-                similarity_output = gr.Markdown(label="Similarity Score")
-                similarity_status = gr.Textbox(label="Status", interactive=False)
-        
-        similarity_btn.click(
-            fn=compute_single_similarity,
-            inputs=[text1_input, text2_input],
-            outputs=[similarity_output, similarity_status]
-        )
 
-# Launch the app
+# =========================
+# Gradio Interface
+# =========================
+SAMPLE_QUERIES = [
+    "dm t2",
+    "छाती में दर्द",
+    "talenovu",
+    "వాంతులు"
+]
+
+with gr.Blocks(title="Parrotlet-E Retrieval (Pinecone)", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🦜 Parrotlet-E: Indic Medical Entities Retrieval
+    Retrieve top-5 semantically similar medical entities from Pinecone using Parrotlet-E embeddings.
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            query_input = gr.Textbox(
+                label="Enter a medical term",
+                placeholder="Type your query here...",
+                lines=1,
+            )
+
+            examples = gr.Examples(
+                examples=SAMPLE_QUERIES,
+                inputs=query_input,
+                label="Example Queries"
+            )
+
+            top_k_slider = gr.Slider(
+                minimum=1,
+                maximum=10,
+                value=5,
+                step=1,
+                label="Top-K Retrieved"
+            )
+
+            search_btn = gr.Button("🔍 Retrieve", variant="primary")
+
+        with gr.Column(scale=2):
+            results_output = gr.Dataframe(
+                headers=["Rank", "Score", "Document", "ID"],
+                datatype=["number", "str", "str", "str"],
+                interactive=False,
+                wrap=True
+            )
+            status_box = gr.Textbox(label="Status", interactive=False)
+
+    search_btn.click(
+        fn=retrieve_documents,
+        inputs=[query_input, top_k_slider],
+        outputs=[results_output, status_box],
+    )
+
+
+# =========================
+# Run App
+# =========================
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860)