Update app.py

app.py

@@ -108,7 +108,7 @@ def compute_similarity(embedder: QwenEmbedder, text1: str, text2: str, model_cho
     similarity = torch.cosine_similarity(embeddings[0:1], embeddings[1:2]).item()
     return round(similarity, 3)
 
-def rerank_documents(embedder: QwenEmbedder, query: str, documents: str) -> List[Tuple[str, float]]:
+def rerank_documents(embedder: QwenEmbedder, query: str, documents: str, model_choice: str = None, embedding_dim: int = None) -> List[Tuple[str, float]]:
     docs_list = [doc.strip() for doc in documents.split('\n') if doc.strip()]
     
     # Add instruction to query
@@ -126,7 +126,7 @@ def rerank_documents(embedder: QwenEmbedder, query: str, documents: str) -> List
     
     return [(doc, round(score, 3)) for doc, score in results]
 
-def process_batch_embeddings(embedder: QwenEmbedder, texts: str) -> pd.DataFrame:
+def process_batch_embeddings(embedder: QwenEmbedder, texts: str, model_choice: str = None, embedding_dim: int = None) -> pd.DataFrame:
     text_list = [text.strip() for text in texts.split('\n') if text.strip()]
     if len(text_list) < 1:
         return pd.DataFrame()
@@ -143,7 +143,7 @@ def process_batch_embeddings(embedder: QwenEmbedder, texts: str) -> pd.DataFrame
     
     return df_similarities.round(3)
 
-def process_retrieval(embedder: QwenEmbedder, task_prompt: str, queries: str, documents: str) -> pd.DataFrame:
+def process_retrieval(embedder: QwenEmbedder, task_prompt: str, queries: str, documents: str, model_choice: str = None, embedding_dim: int = None) -> pd.DataFrame:
     # Process queries and documents
     query_list = [q.strip() for q in queries.split('\n') if q.strip()]
     doc_list = [d.strip() for d in documents.split('\n') if d.strip()]
@@ -165,13 +165,13 @@ def process_retrieval(embedder: QwenEmbedder, task_prompt: str, queries: str, do
     df = pd.DataFrame(scores, index=query_list, columns=doc_list)
     return df.round(3)
 
-def process_cross_lingual(embedder: QwenEmbedder, arabic_text: str, english_text: str) -> dict:
+def process_cross_lingual(embedder: QwenEmbedder, arabic_text: str, english_text: str, model_choice: str = None, embedding_dim: int = None) -> dict:
     texts = [arabic_text, english_text]
     embeddings = embedder.get_embeddings(texts)
     similarity = torch.cosine_similarity(embeddings[0:1], embeddings[1:2]).item()
     return {"similarity": round(similarity, 3)}
 
-def classify_text(embedder: QwenEmbedder, text: str, categories: str) -> List[Tuple[str, float]]:
+def classify_text(embedder: QwenEmbedder, text: str, categories: str, model_choice: str = None, embedding_dim: int = None) -> List[Tuple[str, float]]:
     cat_list = [c.strip() for c in categories.split('\n') if c.strip()]
     text_embedding = embedder.get_embeddings([text])
     cat_embeddings = embedder.get_embeddings(cat_list)
@@ -180,7 +180,7 @@ def classify_text(embedder: QwenEmbedder, text: str, categories: str) -> List[Tu
     results.sort(key=lambda x: x[1], reverse=True)
     return [(cat, round(score, 3)) for cat, score in results]
 
-def cluster_documents(embedder: QwenEmbedder, documents: str, num_clusters: int) -> pd.DataFrame:
+def cluster_documents(embedder: QwenEmbedder, documents: str, num_clusters: int, model_choice: str = None, embedding_dim: int = None) -> pd.DataFrame:
     from sklearn.cluster import KMeans
     doc_list = [doc.strip() for doc in documents.split('\n') if doc.strip()]
     if len(doc_list) < num_clusters:
@@ -212,7 +212,7 @@ def cluster_documents(embedder: QwenEmbedder, documents: str, num_clusters: int)
     })
     return df.sort_values('Cluster')
 
-def analyze_sentiment(embedder: QwenEmbedder, text: str) -> Tuple[str, dict]:
+def analyze_sentiment(embedder: QwenEmbedder, text: str, model_choice: str = None, embedding_dim: int = None) -> Tuple[str, dict]:
     # Define sentiment anchors
     anchors = {
         "very_positive": "هذا رائع جداً ومدهش! أنا سعيد للغاية",
@@ -237,7 +237,7 @@ def analyze_sentiment(embedder: QwenEmbedder, text: str) -> Tuple[str, dict]:
         {k: round(float(v), 3) for k, v in results}
     )
 
-def extract_concepts(embedder: QwenEmbedder, text: str, concept_type: str) -> List[Tuple[str, float]]:
+def extract_concepts(embedder: QwenEmbedder, text: str, concept_type: str, model_choice: str = None, embedding_dim: int = None) -> List[Tuple[str, float]]:
     # Define concept anchors based on type
     concept_anchors = {
         "emotions": [
@@ -693,7 +693,7 @@ def create_demo():
                             similarity_score = gr.Number(label="Similarity Score")
                         
                         similarity_btn.click(
-                            fn=lambda t1, t2, m, d: process_with_embedder('compute_similarity', t1, t2),
+                            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
                         )