Create app.py

app.py

@@ -0,0 +1,201 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModel
+import numpy as np
+from typing import List, Tuple
+import pandas as pd
+
+# Model configuration
+MODEL_NAME = "ekacare/parrotlet-e"
+
+class ParrotletEmbedder:
+    def __init__(self, model_name: str):
+        """Initialize the Parrotlet-E model and tokenizer."""
+        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)
+        self.model.to(self.device)
+        self.model.eval()
+        
+    def mean_pooling(self, model_output, attention_mask):
+        """Perform mean pooling on model output."""
+        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 = []
+        
+        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
+    """
+    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."
+    
+    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
+        
+    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
+if __name__ == "__main__":
+    demo.launch(share=True)