Update app.py

app.py

@@ -12,17 +12,52 @@ from transformers import AutoTokenizer, AutoModel
 # Initialize the embedder at module level
 embedder = None
 
+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
+        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
+
 @spaces.GPU(duration=120)
 def initialize_embedder(embedding_dim=768):
+    # Initialize device inside the GPU worker
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Initializing embedder on device: {device}")
+    
+    # Create and move model to device
+    model = QwenEmbedder(embedding_dim=embedding_dim)
+    return model.to_device(device)
+
+@spaces.GPU(duration=120)
+def process_with_embedder(fn_name, *args):
+    """Generic handler for embedder operations"""
     global embedder
     if embedder is None:
-        # Check for GPU support and configure appropriately
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Initializing embedder on device: {device}")
-        
-        embedder = QwenEmbedder(embedding_dim=embedding_dim)
-        embedder.model = embedder.model.to(device)
-    return embedder
+        embedder = initialize_embedder()
+    
+    # Map function names to actual functions
+    fn_map = {
+        'compute_similarity': compute_similarity,
+        'rerank_documents': rerank_documents,
+        'process_batch_embeddings': process_batch_embeddings,
+        'process_retrieval': process_retrieval,
+        'process_cross_lingual': process_cross_lingual,
+        'classify_text': classify_text,
+        'cluster_documents': cluster_documents,
+        'analyze_sentiment': analyze_sentiment,
+        'extract_concepts': extract_concepts
+    }
+    
+    return fn_map[fn_name](embedder, *args)
 
 # Check for GPU support and configure appropriately
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -49,40 +84,6 @@ def tokenize(tokenizer, input_texts, eod_id, max_length):
     batch_dict = tokenizer.pad(batch_dict, padding=True, return_tensors="pt")
     return batch_dict
 
-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')
-        # Uncomment below for better performance if GPU available
-        # self.model = AutoModel.from_pretrained('Qwen/Qwen3-Embedding-0.6B', 
-        #     attn_implementation="flash_attention_2", 
-        #     torch_dtype=torch.float16
-        # ).cuda()
-        self.eod_id = self.tokenizer.convert_tokens_to_ids("<|endoftext|>")
-        self.max_length = 8192
-        self.embedding_dim = embedding_dim
-        self.projection = torch.nn.Linear(768, embedding_dim) if embedding_dim != 768 else None
-        
-    def get_embeddings(self, texts: List[str], with_instruction: bool = False) -> Tensor:
-        if with_instruction:
-            task = 'Process and understand the following text'
-            texts = [get_detailed_instruct(task, text) for text in texts]
-        
-        batch_dict = tokenize(self.tokenizer, texts, self.eod_id, self.max_length)
-        batch_dict.to(self.model.device)
-        
-        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)
-        
-        return embeddings
-
 def compute_similarity(embedder: QwenEmbedder, text1: str, text2: str) -> float:
     embeddings = embedder.get_embeddings([text1, text2])
     similarity = torch.cosine_similarity(embeddings[0:1], embeddings[1:2]).item()
@@ -259,12 +260,6 @@ def extract_concepts(embedder: QwenEmbedder, text: str, concept_type: str) -> Li
     
     return [(concept, round(score, 3)) for concept, score in results]
 
-# Add a function to reinitialize embedder with new dimension
-def reinitialize_embedder(dim: int) -> QwenEmbedder:
-    global embedder
-    embedder = QwenEmbedder(embedding_dim=dim)
-    return "Embedder reinitialized with dimension: " + str(dim)
-
 # Update the CSS to improve feature visibility
 custom_css = """
 :root {
@@ -454,16 +449,9 @@ button.secondary {
 
 # Create the Gradio interface
 def create_demo():
-    global embedder
-    # Initialize embedder with GPU support
-    embedder = initialize_embedder()
-    
     demo = gr.Blocks(title="Advanced Text Processing with Qwen", css=custom_css, theme=gr.themes.Soft())
     
     with demo:
-        # Store embedder in state
-        state = gr.State(embedder)
-        
         with gr.Row():
             # Sidebar
             with gr.Column(scale=1, elem_classes="sidebar"):
@@ -601,8 +589,8 @@ def create_demo():
                             similarity_score = gr.Number(label="Similarity Score")
                         
                         similarity_btn.click(
-                            fn=lambda t1, t2, s: compute_similarity(s.value, t1, t2),
-                            inputs=[text1, text2, state],
+                            fn=lambda t1, t2: process_with_embedder('compute_similarity', t1, t2),
+                            inputs=[text1, text2],
                             outputs=similarity_score
                         )
                     
@@ -644,8 +632,8 @@ def create_demo():
                             )
                         
                         rerank_btn.click(
-                            fn=lambda q, d, s: rerank_documents(s.value, q, d),
-                            inputs=[query_text, documents_text, state],
+                            fn=lambda q, d: process_with_embedder('rerank_documents', q, d),
+                            inputs=[query_text, documents_text],
                             outputs=rerank_results
                         )
                     
@@ -679,8 +667,8 @@ def create_demo():
                             )
                         
                         process_btn.click(
-                            fn=lambda t, s: process_batch_embeddings(s.value, t),
-                            inputs=[batch_texts, state],
+                            fn=lambda t: process_with_embedder('process_batch_embeddings', t),
+                            inputs=[batch_texts],
                             outputs=[similarity_matrix]
                         )
 
@@ -751,8 +739,8 @@ def create_demo():
                             """)
                         
                         retrieve_btn.click(
-                            fn=lambda p, q, d, s: process_retrieval(s.value, p, q, d),
-                            inputs=[task_prompt, queries_text, documents_text, state],
+                            fn=lambda p, q, d: process_with_embedder('process_retrieval', p, q, d),
+                            inputs=[task_prompt, queries_text, documents_text],
                             outputs=[retrieval_matrix]
                         )
 
@@ -807,8 +795,8 @@ def create_demo():
                             """)
                         
                         match_btn.click(
-                            fn=lambda a, e, s: process_cross_lingual(s.value, a, e)["similarity"],
-                            inputs=[arabic_text, english_text, state],
+                            fn=lambda a, e: process_with_embedder('process_cross_lingual', a, e),
+                            inputs=[arabic_text, english_text],
                             outputs=[cross_lingual_score]
                         )
 
@@ -850,8 +838,8 @@ def create_demo():
                             )
                         
                         classify_btn.click(
-                            fn=lambda t, c, s: classify_text(s.value, t, c),
-                            inputs=[input_text, categories_text, state],
+                            fn=lambda t, c: process_with_embedder('classify_text', t, c),
+                            inputs=[input_text, categories_text],
                             outputs=classification_results
                         )
 
@@ -899,8 +887,8 @@ def create_demo():
                             )
                         
                         cluster_btn.click(
-                            fn=lambda d, n, s: cluster_documents(s.value, d, n),
-                            inputs=[cluster_docs, num_clusters, state],
+                            fn=lambda d, n: process_with_embedder('cluster_documents', d, n),
+                            inputs=[cluster_docs, num_clusters],
                             outputs=clustering_results
                         )
 
@@ -932,8 +920,8 @@ def create_demo():
                                 sentiment_scores = gr.Json(label="Detailed Scores")
                             
                             analyze_btn.click(
-                                fn=lambda t, s: analyze_sentiment(s.value, t),
-                                inputs=[sentiment_text, state],
+                                fn=lambda t: process_with_embedder('analyze_sentiment', t),
+                                inputs=[sentiment_text],
                                 outputs=[sentiment_label, sentiment_scores]
                             )
 
@@ -972,29 +960,30 @@ def create_demo():
                             )
                             
                             extract_btn.click(
-                                fn=lambda t, c, s: extract_concepts(s.value, t, c),
-                                inputs=[concept_text, concept_type, state],
+                                fn=lambda t, c: process_with_embedder('extract_concepts', t, c),
+                                inputs=[concept_text, concept_type],
                                 outputs=concept_results
                             )
 
-        # Fix dimension update functionality
-        def update_embedder_dim(dim, state):
+        # Update dimension handler
+        @spaces.GPU(duration=120)
+        def update_embedder_dim(dim):
+            global embedder
             try:
-                new_embedder = initialize_embedder(embedding_dim=dim)
-                state.value = new_embedder
-                return state, f"Successfully updated embedding dimension to {dim}"
+                embedder = initialize_embedder(embedding_dim=dim)
+                return f"Successfully updated embedding dimension to {dim}"
             except Exception as e:
-                return state, f"Error updating dimension: {str(e)}"
+                return f"Error updating dimension: {str(e)}"
         
         update_dim_btn.click(
             fn=update_embedder_dim,
-            inputs=[embedding_dim, state],
-            outputs=[state, dim_status]
+            inputs=[embedding_dim],
+            outputs=dim_status
         )
     
     return demo
 
 if __name__ == "__main__":
     demo = create_demo()
-    demo.queue()  # Enable queuing for better handling of GPU resources
+    demo.queue()
     demo.launch()