Update app.py

app.py

@@ -5,46 +5,45 @@ import numpy as np
 import plotly.express as px
 import pandas as pd
 import spaces
-from typing import List, Tuple
+from typing import List, Tuple, Dict
 from torch import Tensor
 from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import SentenceTransformer
+import json
 
 # Initialize the embedder at module level
 embedder = None
 
+AVAILABLE_MODELS = {
+    "Qwen Original": "Qwen/Qwen3-Embedding-0.6B",
+    "Arabic Fine-tuned": "Omartificial-Intelligence-Space/Semantic-Ar-Qwen-Embed-0.6B"
+}
+
 class QwenEmbedder:
-    def __init__(self, embedding_dim=768):
-        self.tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B', padding_side='left')
-        self.model = AutoModel.from_pretrained('Qwen/Qwen3-Embedding-0.6B')
-        self.eod_id = self.tokenizer.convert_tokens_to_ids("<|endoftext|>")
-        self.max_length = 8192
+    def __init__(self, model_name: str = "Qwen/Qwen3-Embedding-0.6B", embedding_dim: int = 768):
+        self.model = SentenceTransformer(model_name)
         self.embedding_dim = embedding_dim
-        self.projection = torch.nn.Linear(768, embedding_dim) if embedding_dim != 768 else None
-        
-    def to_device(self, device):
-        self.model = self.model.to(device)
-        if self.projection is not None:
-            self.projection = self.projection.to(device)
-        return self
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
         
-    def get_embeddings(self, texts: List[str], with_instruction: bool = False) -> Tensor:
+        if embedding_dim != 768:
+            # Add projection layer if needed
+            self.projection = torch.nn.Linear(768, embedding_dim)
+            self.projection.to(self.device)
+        else:
+            self.projection = None
+
+    def get_embeddings(self, texts: List[str], with_instruction: bool = False) -> torch.Tensor:
         if with_instruction:
-            task = 'Process and understand the following text'
-            texts = [get_detailed_instruct(task, text) for text in texts]
+            texts = [f"Represent this Arabic text for retrieval: {text}" for text in texts]
         
-        batch_dict = tokenize(self.tokenizer, texts, self.eod_id, self.max_length)
-        batch_dict = {k: v.to(self.model.device) for k, v in batch_dict.items()}
+        embeddings = self.model.encode(texts, convert_to_tensor=True)
         
-        with torch.no_grad():
-            outputs = self.model(**batch_dict)
-            embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
-            
-            # Project to desired dimension if needed
-            if self.projection is not None:
-                embeddings = self.projection(embeddings)
-                
-            embeddings = F.normalize(embeddings, p=2, dim=1)
+        if self.projection is not None:
+            embeddings = self.projection(embeddings)
         
+        # Normalize embeddings
+        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
         return embeddings
 
 @spaces.GPU(duration=120)
@@ -280,6 +279,86 @@ def extract_concepts(embedder: QwenEmbedder, text: str, concept_type: str) -> Li
     
     return [(concept, round(score, 3)) for concept, score in results]
 
+def create_embedder(model_choice: str, embedding_dim: int = 768) -> QwenEmbedder:
+    model_name = AVAILABLE_MODELS[model_choice]
+    return QwenEmbedder(model_name=model_name, embedding_dim=embedding_dim)
+
+def process_similarity(text1: str, text2: str, model_choice: str, embedding_dim: int) -> float:
+    embedder = create_embedder(model_choice, embedding_dim)
+    embeddings = embedder.get_embeddings([text1, text2])
+    similarity = torch.nn.functional.cosine_similarity(embeddings[0].unsqueeze(0), embeddings[1].unsqueeze(0))
+    return float(similarity)
+
+def process_reranking(query: str, documents: str, model_choice: str, embedding_dim: int) -> Dict:
+    embedder = create_embedder(model_choice, embedding_dim)
+    documents = [doc.strip() for doc in documents.split('\n') if doc.strip()]
+    
+    query_embedding = embedder.get_embeddings([query], with_instruction=True)
+    doc_embeddings = embedder.get_embeddings(documents)
+    
+    similarities = torch.nn.functional.cosine_similarity(query_embedding, doc_embeddings)
+    
+    # Sort documents by similarity
+    sorted_indices = torch.argsort(similarities, descending=True)
+    results = []
+    for idx in sorted_indices:
+        results.append({
+            'document': documents[idx],
+            'score': float(similarities[idx])
+        })
+    
+    return {'results': results}
+
+def process_batch(texts: str, model_choice: str, embedding_dim: int) -> Dict:
+    embedder = create_embedder(model_choice, embedding_dim)
+    texts = [text.strip() for text in texts.split('\n') if text.strip()]
+    
+    embeddings = embedder.get_embeddings(texts)
+    similarity_matrix = torch.nn.functional.cosine_similarity(embeddings.unsqueeze(1), embeddings.unsqueeze(0), dim=2)
+    
+    df = pd.DataFrame(similarity_matrix.cpu().numpy(), index=texts, columns=texts)
+    return {'similarity_matrix': df.to_dict()}
+
+def process_retrieval(prompt: str, queries: str, documents: str, model_choice: str, embedding_dim: int) -> Dict:
+    embedder = create_embedder(model_choice, embedding_dim)
+    
+    # Process input strings
+    queries = [q.strip() for q in queries.split('\n') if q.strip()]
+    documents = [doc.strip() for doc in documents.split('\n') if doc.strip()]
+    
+    # Get embeddings
+    prompt_embedding = embedder.get_embeddings([prompt], with_instruction=True)
+    query_embeddings = embedder.get_embeddings(queries, with_instruction=True)
+    doc_embeddings = embedder.get_embeddings(documents)
+    
+    # Calculate similarities
+    query_similarities = torch.nn.functional.cosine_similarity(prompt_embedding, query_embeddings)
+    doc_similarities = torch.nn.functional.cosine_similarity(prompt_embedding.repeat(len(documents), 1), doc_embeddings)
+    
+    # Process results
+    results = {
+        'relevant_queries': [],
+        'relevant_documents': []
+    }
+    
+    # Sort queries
+    query_indices = torch.argsort(query_similarities, descending=True)
+    for idx in query_indices:
+        results['relevant_queries'].append({
+            'query': queries[idx],
+            'similarity': float(query_similarities[idx])
+        })
+    
+    # Sort documents
+    doc_indices = torch.argsort(doc_similarities, descending=True)
+    for idx in doc_indices:
+        results['relevant_documents'].append({
+            'document': documents[idx],
+            'similarity': float(doc_similarities[idx])
+        })
+    
+    return results
+
 # Update the CSS to improve feature visibility
 custom_css = """
 :root {
@@ -490,6 +569,11 @@ def create_demo():
                     Configure the embedding model parameters below.
                     """)
                     
+                    model_choice = gr.Dropdown(
+                        choices=list(AVAILABLE_MODELS.keys()),
+                        value=list(AVAILABLE_MODELS.keys())[0],
+                        label="Select Model"
+                    )
                     embedding_dim = gr.Slider(
                         minimum=32,
                         maximum=1024,
@@ -609,8 +693,8 @@ def create_demo():
                             similarity_score = gr.Number(label="Similarity Score")
                         
                         similarity_btn.click(
-                            fn=lambda t1, t2: process_with_embedder('compute_similarity', t1, t2),
-                            inputs=[text1, text2],
+                            fn=lambda t1, t2, m, d: process_with_embedder('compute_similarity', t1, t2, m, d),
+                            inputs=[text1, text2, model_choice, embedding_dim],
                             outputs=similarity_score
                         )
                     
@@ -652,8 +736,8 @@ def create_demo():
                             )
                         
                         rerank_btn.click(
-                            fn=lambda q, d: process_with_embedder('rerank_documents', q, d),
-                            inputs=[query_text, documents_text],
+                            fn=lambda q, d, m, e: process_with_embedder('rerank_documents', q, d, m, e),
+                            inputs=[query_text, documents_text, model_choice, embedding_dim],
                             outputs=rerank_results
                         )
                     
@@ -687,8 +771,8 @@ def create_demo():
                             )
                         
                         process_btn.click(
-                            fn=lambda t: process_with_embedder('process_batch_embeddings', t),
-                            inputs=[batch_texts],
+                            fn=lambda t, m, e: process_with_embedder('process_batch_embeddings', t, m, e),
+                            inputs=[batch_texts, model_choice, embedding_dim],
                             outputs=[similarity_matrix]
                         )
 
@@ -759,8 +843,8 @@ def create_demo():
                             """)
                         
                         retrieve_btn.click(
-                            fn=lambda p, q, d: process_with_embedder('process_retrieval', p, q, d),
-                            inputs=[task_prompt, queries_text, documents_text],
+                            fn=lambda p, q, d, m, e: process_with_embedder('process_retrieval', p, q, d, m, e),
+                            inputs=[task_prompt, queries_text, documents_text, model_choice, embedding_dim],
                             outputs=[retrieval_matrix]
                         )
 
@@ -815,8 +899,8 @@ def create_demo():
                             """)
                         
                         match_btn.click(
-                            fn=lambda a, e: process_with_embedder('process_cross_lingual', a, e),
-                            inputs=[arabic_text, english_text],
+                            fn=lambda a, e, m, e: process_with_embedder('process_cross_lingual', a, e, m, e),
+                            inputs=[arabic_text, english_text, model_choice, embedding_dim],
                             outputs=[cross_lingual_score]
                         )
 
@@ -858,8 +942,8 @@ def create_demo():
                             )
                         
                         classify_btn.click(
-                            fn=lambda t, c: process_with_embedder('classify_text', t, c),
-                            inputs=[input_text, categories_text],
+                            fn=lambda t, c, m, e: process_with_embedder('classify_text', t, c, m, e),
+                            inputs=[input_text, categories_text, model_choice, embedding_dim],
                             outputs=classification_results
                         )
 
@@ -907,8 +991,8 @@ def create_demo():
                             )
                         
                         cluster_btn.click(
-                            fn=lambda d, n: process_with_embedder('cluster_documents', d, n),
-                            inputs=[cluster_docs, num_clusters],
+                            fn=lambda d, n, m, e: process_with_embedder('cluster_documents', d, n, m, e),
+                            inputs=[cluster_docs, num_clusters, model_choice, embedding_dim],
                             outputs=clustering_results
                         )
 
@@ -940,8 +1024,8 @@ def create_demo():
                                 sentiment_scores = gr.Json(label="Detailed Scores")
                             
                             analyze_btn.click(
-                                fn=lambda t: process_with_embedder('analyze_sentiment', t),
-                                inputs=[sentiment_text],
+                                fn=lambda t, m, e: process_with_embedder('analyze_sentiment', t, m, e),
+                                inputs=[sentiment_text, model_choice, embedding_dim],
                                 outputs=[sentiment_label, sentiment_scores]
                             )
 
@@ -980,8 +1064,8 @@ def create_demo():
                             )
                             
                             extract_btn.click(
-                                fn=lambda t, c: process_with_embedder('extract_concepts', t, c),
-                                inputs=[concept_text, concept_type],
+                                fn=lambda t, c, m, e: process_with_embedder('extract_concepts', t, c, m, e),
+                                inputs=[concept_text, concept_type, model_choice, embedding_dim],
                                 outputs=concept_results
                             )