Spaces:
Running
Running
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoModel | |
| import torch | |
| # --- Model Loading --- | |
| tokenizer_splade = None | |
| model_splade = None | |
| tokenizer_splade_lexical = None | |
| model_splade_lexical = None | |
| tokenizer_splade_doc = None | |
| model_splade_doc = None | |
| # Load SPLADE v3 model (original) | |
| try: | |
| tokenizer_splade = AutoTokenizer.from_pretrained("naver/splade-cocondenser-selfdistil") | |
| model_splade = AutoModelForMaskedLM.from_pretrained("naver/splade-cocondenser-selfdistil") | |
| model_splade.eval() # Set to evaluation mode for inference | |
| print("SPLADE-cocondenser-distil model loaded successfully!") | |
| except Exception as e: | |
| print(f"Error loading SPLADE-cocondenser-distil model: {e}") | |
| print("Please ensure you have accepted any user access agreements on the Hugging Face Hub page for 'naver/splade-cocondenser-selfdistil'.") | |
| # Load SPLADE v3 Lexical model | |
| try: | |
| splade_lexical_model_name = "naver/splade-v3-lexical" | |
| tokenizer_splade_lexical = AutoTokenizer.from_pretrained(splade_lexical_model_name) | |
| model_splade_lexical = AutoModelForMaskedLM.from_pretrained(splade_lexical_model_name) | |
| model_splade_lexical.eval() # Set to evaluation mode for inference | |
| print(f"SPLADE-v3-Lexical model '{splade_lexical_model_name}' loaded successfully!") | |
| except Exception as e: | |
| print(f"Error loading SPLADE-v3-Lexical model: {e}") | |
| print(f"Please ensure '{splade_lexical_model_name}' is accessible (check Hugging Face Hub for potential agreements).") | |
| # Load SPLADE v3 Doc model | |
| try: | |
| splade_doc_model_name = "naver/splade-v3-doc" | |
| tokenizer_splade_doc = AutoTokenizer.from_pretrained(splade_doc_model_name) | |
| model_splade_doc = AutoModelForMaskedLM.from_pretrained(splade_doc_model_name) | |
| model_splade_doc.eval() # Set to evaluation mode for inference | |
| print(f"SPLADE-v3-Doc model '{splade_doc_model_name}' loaded successfully!") | |
| except Exception as e: | |
| print(f"Error loading SPLADE-v3-Doc model: {e}") | |
| print(f"Please ensure '{splade_doc_model_name}' is accessible (check Hugging Face Hub for potential agreements).") | |
| # --- Helper function for lexical mask --- | |
| def create_lexical_bow_mask(input_ids, vocab_size, tokenizer): | |
| """ | |
| Creates a binary bag-of-words mask from input_ids, | |
| zeroing out special tokens and padding. | |
| """ | |
| bow_mask = torch.zeros(vocab_size, device=input_ids.device) | |
| meaningful_token_ids = [] | |
| for token_id in input_ids.squeeze().tolist(): | |
| if token_id not in [ | |
| tokenizer.pad_token_id, | |
| tokenizer.cls_token_id, | |
| tokenizer.sep_token_id, | |
| tokenizer.mask_token_id, | |
| tokenizer.unk_token_id | |
| ]: | |
| meaningful_token_ids.append(token_id) | |
| if meaningful_token_ids: | |
| bow_mask[list(set(meaningful_token_ids))] = 1 | |
| return bow_mask.unsqueeze(0) | |
| # --- Core Representation Functions --- | |
| def get_splade_cocondenser_representation(text): | |
| if tokenizer_splade is None or model_splade is None: | |
| return "SPLADE-cocondenser-distil model is not loaded. Please check the console for loading errors." | |
| inputs = tokenizer_splade(text, return_tensors="pt", padding=True, truncation=True) | |
| inputs = {k: v.to(model_splade.device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| output = model_splade(**inputs) | |
| if hasattr(output, 'logits'): | |
| # Standard SPLADE calculation for learned weighting and expansion | |
| splade_vector = torch.max( | |
| torch.log(1 + torch.relu(output.logits)) * inputs['attention_mask'].unsqueeze(-1), | |
| dim=1 | |
| )[0].squeeze() | |
| else: | |
| return "Model output structure not as expected for SPLADE-cocondenser-distil. 'logits' not found." | |
| indices = torch.nonzero(splade_vector).squeeze().cpu().tolist() | |
| if not isinstance(indices, list): | |
| indices = [indices] | |
| values = splade_vector[indices].cpu().tolist() | |
| token_weights = dict(zip(indices, values)) | |
| meaningful_tokens = {} | |
| for token_id, weight in token_weights.items(): | |
| decoded_token = tokenizer_splade.decode([token_id]) | |
| if decoded_token not in ["[CLS]", "[SEP]", "[PAD]", "[UNK]"] and len(decoded_token.strip()) > 0: | |
| meaningful_tokens[decoded_token] = weight | |
| sorted_representation = sorted(meaningful_tokens.items(), key=lambda item: item[1], reverse=True) | |
| formatted_output = "SPLADE-cocondenser-distil Representation (Weighting and Expansion):\n" | |
| if not sorted_representation: | |
| formatted_output += "No significant terms found for this input.\n" | |
| else: | |
| for term, weight in sorted_representation: | |
| formatted_output += f"- **{term}**: {weight:.4f}\n" | |
| formatted_output += "\n--- Raw SPLADE Vector Info ---\n" | |
| formatted_output += f"Total non-zero terms in vector: {len(indices)}\n" | |
| formatted_output += f"Sparsity: {1 - (len(indices) / tokenizer_splade.vocab_size):.2%}\n" | |
| return formatted_output | |
| def get_splade_lexical_representation(text): | |
| 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." | |
| inputs = tokenizer_splade_lexical(text, return_tensors="pt", padding=True, truncation=True) | |
| inputs = {k: v.to(model_splade_lexical.device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| output = model_splade_lexical(**inputs) | |
| if hasattr(output, 'logits'): | |
| splade_vector = torch.max( | |
| torch.log(1 + torch.relu(output.logits)) * inputs['attention_mask'].unsqueeze(-1), | |
| dim=1 | |
| )[0].squeeze() | |
| else: | |
| return "Model output structure not as expected for SPLADE-v3-Lexical. 'logits' not found." | |
| # Always apply lexical mask for this model's specific behavior | |
| vocab_size = tokenizer_splade_lexical.vocab_size | |
| bow_mask = create_lexical_bow_mask( | |
| inputs['input_ids'], vocab_size, tokenizer_splade_lexical | |
| ).squeeze() | |
| splade_vector = splade_vector * bow_mask | |
| indices = torch.nonzero(splade_vector).squeeze().cpu().tolist() | |
| if not isinstance(indices, list): | |
| indices = [indices] | |
| values = splade_vector[indices].cpu().tolist() | |
| token_weights = dict(zip(indices, values)) | |
| meaningful_tokens = {} | |
| for token_id, weight in token_weights.items(): | |
| decoded_token = tokenizer_splade_lexical.decode([token_id]) | |
| if decoded_token not in ["[CLS]", "[SEP]", "[PAD]", "[UNK]"] and len(decoded_token.strip()) > 0: | |
| meaningful_tokens[decoded_token] = weight | |
| sorted_representation = sorted(meaningful_tokens.items(), key=lambda item: item[1], reverse=True) | |
| formatted_output = "SPLADE-v3-Lexical Representation (Weighting):\n" | |
| if not sorted_representation: | |
| formatted_output += "No significant terms found for this input.\n" | |
| else: | |
| for term, weight in sorted_representation: | |
| formatted_output += f"- **{term}**: {weight:.4f}\n" | |
| formatted_output += "\n--- Raw SPLADE Vector Info ---\n" | |
| formatted_output += f"Total non-zero terms in vector: {len(indices)}\n" | |
| formatted_output += f"Sparsity: {1 - (len(indices) / tokenizer_splade_lexical.vocab_size):.2%}\n" | |
| return formatted_output | |
| # Function for SPLADE-v3-Doc representation (Binary Sparse - Lexical Only) | |
| def get_splade_doc_representation(text): | |
| if tokenizer_splade_doc is None or model_splade_doc is None: | |
| return "SPLADE-v3-Doc model is not loaded. Please check the console for loading errors." | |
| inputs = tokenizer_splade_doc(text, return_tensors="pt", padding=True, truncation=True) | |
| inputs = {k: v.to(model_splade_doc.device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| output = model_splade_doc(**inputs) | |
| if not hasattr(output, "logits"): | |
| return "SPLADE-v3-Doc model output structure not as expected. 'logits' not found." | |
| # For SPLADE-v3-Doc, assuming output is designed to be binary and lexical-only. | |
| # We will derive the output directly from the input tokens themselves, | |
| # as the model's primary role in this context is as a pre-trained LM feature extractor | |
| # for a document-side, lexical-only binary sparse representation. | |
| vocab_size = tokenizer_splade_doc.vocab_size | |
| binary_splade_vector = create_lexical_bow_mask( # Use the BOW mask directly for binary | |
| inputs['input_ids'], vocab_size, tokenizer_splade_doc | |
| ).squeeze() | |
| indices = torch.nonzero(binary_splade_vector).squeeze().cpu().tolist() | |
| if not isinstance(indices, list): | |
| indices = [indices] if indices else [] | |
| values = [1.0] * len(indices) # All values are 1 for binary representation | |
| token_weights = dict(zip(indices, values)) | |
| meaningful_tokens = {} | |
| for token_id, weight in token_weights.items(): | |
| decoded_token = tokenizer_splade_doc.decode([token_id]) | |
| if decoded_token not in ["[CLS]", "[SEP]", "[PAD]", "[UNK]"] and len(decoded_token.strip()) > 0: | |
| meaningful_tokens[decoded_token] = weight | |
| sorted_representation = sorted(meaningful_tokens.items(), key=lambda item: item[0]) # Sort alphabetically for clarity | |
| formatted_output = "SPLADE-v3-Doc Representation (Binary):\n" | |
| if not sorted_representation: | |
| formatted_output += "No significant terms found for this input.\n" | |
| else: | |
| for i, (term, _) in enumerate(sorted_representation): | |
| if i >= 50: # Limit display for very long lists | |
| formatted_output += f"...and {len(sorted_representation) - 50} more terms.\n" | |
| break | |
| formatted_output += f"- **{term}**\n" | |
| formatted_output += "\n--- Raw Binary Sparse Vector Info ---\n" | |
| formatted_output += f"Total activated terms: {len(indices)}\n" | |
| formatted_output += f"Sparsity: {1 - (len(indices) / tokenizer_splade_doc.vocab_size):.2%}\n" | |
| return formatted_output | |
| # --- Unified Prediction Function for Gradio --- | |
| def predict_representation(model_choice, text): | |
| if model_choice == "SPLADE-cocondenser-distil (weighting and expansion)": | |
| return get_splade_cocondenser_representation(text) | |
| elif model_choice == "SPLADE-v3-Lexical (weighting)": | |
| return get_splade_lexical_representation(text) | |
| elif model_choice == "SPLADE-v3-Doc (binary)": | |
| return get_splade_doc_representation(text) | |
| else: | |
| return "Please select a model." | |
| # --- Gradio Interface Setup --- | |
| demo = gr.Interface( | |
| fn=predict_representation, | |
| inputs=[ | |
| gr.Radio( | |
| [ | |
| "SPLADE-cocondenser-distil (weighting and expansion)", | |
| "SPLADE-v3-Lexical (weighting)", | |
| "SPLADE-v3-Doc (binary)" | |
| ], | |
| label="Choose Representation Model", | |
| value="SPLADE-cocondenser-distil (weighting and expansion)" # Corrected default value | |
| ), | |
| gr.Textbox( | |
| lines=5, | |
| label="Enter your query or document text here:", | |
| placeholder="e.g., Why is Padua the nicest city in Italy?" | |
| ) | |
| ], | |
| outputs=gr.Markdown(), | |
| title="🌌 Sparse Representation Generator", | |
| description="Explore different SPLADE models and their sparse representation types: weighted and expansive, weighted and lexical-only, or strictly binary.", | |
| allow_flagging="never" | |
| ) | |
| # Launch the Gradio app | |
| demo.launch() |