Update app.py

app.py

@@ -107,7 +107,7 @@ def get_splade_representation(text):
     return formatted_output
 
 
-def get_splade_lexical_representation(text, apply_lexical_mask: bool): # Added parameter
+def get_splade_lexical_representation(text): # Removed apply_lexical_mask parameter
     if tokenizer_splade_lexical is None or model_splade_lexical is None:
         return "SPLADE v3 Lexical model is not loaded. Please check the console for loading errors."
 
@@ -125,18 +125,17 @@ def get_splade_lexical_representation(text, apply_lexical_mask: bool): # Added p
     else:
         return "Model output structure not as expected for SPLADE v3 Lexical. 'logits' not found."
 
-    # --- Apply Lexical Mask if requested ---
-    if apply_lexical_mask:
-        # Get the vocabulary size from the tokenizer
-        vocab_size = tokenizer_splade_lexical.vocab_size
-        
-        # Create the Bag-of-Words mask
-        bow_mask = create_lexical_bow_mask(
-            inputs['input_ids'], vocab_size, tokenizer_splade_lexical
-        ).squeeze() # Squeeze to match splade_vector's [vocab_size] shape
-        
-        # Multiply the SPLADE vector by the BoW mask to zero out expanded terms
-        splade_vector = splade_vector * bow_mask
+    # --- Apply Lexical Mask (always applied for this function now) ---
+    # Get the vocabulary size from the tokenizer
+    vocab_size = tokenizer_splade_lexical.vocab_size
+    
+    # Create the Bag-of-Words mask
+    bow_mask = create_lexical_bow_mask(
+        inputs['input_ids'], vocab_size, tokenizer_splade_lexical
+    ).squeeze() 
+    
+    # Multiply the SPLADE vector by the BoW mask to zero out expanded terms
+    splade_vector = splade_vector * bow_mask
     # --- End Lexical Mask Logic ---
 
     indices = torch.nonzero(splade_vector).squeeze().cpu().tolist()
@@ -172,12 +171,8 @@ def get_splade_lexical_representation(text, apply_lexical_mask: bool): # Added p
 def predict_representation(model_choice, text):
     if model_choice == "SPLADE (cocondenser)":
         return get_splade_representation(text)
-    elif model_choice == "SPLADE-v3-Lexical (with expansion)":
-        # Call the lexical function without applying the mask
-        return get_splade_lexical_representation(text, apply_lexical_mask=False)
-    elif model_choice == "SPLADE-v3-Lexical (lexical-only)":
-        # Call the lexical function applying the mask
-        return get_splade_lexical_representation(text, apply_lexical_mask=True)
+    elif model_choice == "SPLADE-v3-Lexical": # Simplified choice
+        return get_splade_lexical_representation(text) # Always applies lexical mask
     else:
         return "Please select a model."
 
@@ -188,8 +183,7 @@ demo = gr.Interface(
         gr.Radio(
             [
                 "SPLADE (cocondenser)",
-                "SPLADE-v3-Lexical (with expansion)", # Option to see full neural output
-                "SPLADE-v3-Lexical (lexical-only)"    # Option with lexical mask applied
+                "SPLADE-v3-Lexical" # Simplified option
             ],
             label="Choose Representation Model",
             value="SPLADE (cocondenser)" # Default selection
@@ -202,7 +196,7 @@ demo = gr.Interface(
     ],
     outputs=gr.Markdown(),
     title="🌌 Sparse Representation Generator",
-    description="Enter any text to see its SPLADE sparse vector. Explore the difference between full neural expansion and lexical-only representations.",
+    description="Enter any text to see its SPLADE sparse vector.", # Simplified description
     allow_flagging="never"
 )