initial

app.py

@@ -1,7 +1,139 @@
+from typing import List, Dict, Any
+
 import gradio as gr
+import spaces
+import torch
+import numpy as np
+
+# For the dense embedding
+from sentence_transformers import SentenceTransformer
+
+# For SPLADE sparse embedding
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+# For ColBERT
+from transformers import AutoModel, AutoTokenizer
+
+
+############################
+# 1) Load models & tokenizers
+############################
+
+# 1A) Dense embedding model (Nomic)
+dense_model = SentenceTransformer(
+    "nomic-ai/nomic-embed-text-v1.5",
+    trust_remote_code=True,
+    device="cuda"  # Force GPU if available
+)
+
+# 1B) SPLADE for sparse embeddings
+#    Using "naver/splade-cocondenser-ensembledistil" as an example
+sparse_tokenizer = AutoTokenizer.from_pretrained("naver/splade-cocondenser-ensembledistil")
+sparse_model = AutoModelForMaskedLM.from_pretrained("naver/splade-cocondenser-ensembledistil")
+sparse_model.eval()
+sparse_model.to("cuda")  # move to GPU
+
+# 1C) ColBERT model
+colbert_tokenizer = AutoTokenizer.from_pretrained("colbert-ir/colbertv2.0")
+colbert_model = AutoModel.from_pretrained("colbert-ir/colbertv2.0")
+colbert_model.eval()
+colbert_model.to("cuda")
+
+
+############################
+# 2) Helper functions
+############################
+
+def get_dense_embedding(text: str) -> List[float]:
+    """
+    Use SentenceTransformer to get a single dense vector.
+    """
+    # model.encode returns a NumPy array of shape (dim,)
+    emb = dense_model.encode(text)
+    return emb.tolist()  # convert to Python list for JSON serialization
+
+
+def get_splade_sparse_embedding(text: str) -> List[float]:
+    """
+    Compute a sparse embedding with SPLADE (max pooling over tokens).
+    Returns a large vector ~ vocabulary size, e.g. 30k+ dims.
+    """
+    inputs = sparse_tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=256
+    )
+    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
+    with torch.no_grad():
+        # shape: [batch=1, seq_len, vocab_size]
+        logits = sparse_model(**inputs).logits.squeeze(0)  # [seq_len, vocab_size]
+
+    # SPLADE approach for query-like encoding (max over sequence dimension):
+    # For doc encoding, one might do sum instead of max; usage can differ.
+    # We'll do max pooling: log(1 + ReLU(logits)) -> max over seq_len
+    sparse_emb = torch.log1p(torch.relu(logits)).max(dim=0).values
+    # Convert to CPU list
+    return sparse_emb.cpu().numpy().tolist()
+
+
+def get_colbert_embedding(text: str) -> List[List[float]]:
+    """
+    Generate token-level embeddings via ColBERT.
+    Returns a list of [token_dim] for each token in the sequence.
+    """
+    inputs = colbert_tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=180
+    )
+    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
+    with torch.no_grad():
+        outputs = colbert_model(**inputs)
+        # outputs.last_hidden_state: [1, seq_len, hidden_dim]
+        emb = outputs.last_hidden_state.squeeze(0)  # shape: [seq_len, hidden_dim]
+
+    # Convert each token embedding to a list
+    return emb.cpu().numpy().tolist()
+
+
+############################
+# 3) The main embedding function
+############################
+
+@spaces.GPU
+def embed(document: str) -> Dict[str, Any]:
+    """
+    Single function that returns dense, sparse (SPLADE), and ColBERT embeddings.
+    Decorated with @spaces.GPU for ephemeral GPU usage in Hugging Face Spaces.
+    """
+    dense_emb = get_dense_embedding(document)
+    sparse_emb = get_splade_sparse_embedding(document)
+    colbert_emb = get_colbert_embedding(document)
+
+    return {
+        "dense_embedding": dense_emb,
+        "sparse_embedding": sparse_emb,
+        "colbert_embedding": colbert_emb
+    }
+
+
+############################
+# 4) Gradio App
+############################
+
+with gr.Blocks() as app:
+    gr.Markdown("# Multi-Embedding Generator (Dense, SPLADE, ColBERT)")
+
+    text_input = gr.Textbox(label="Enter text to embed")
+    output = gr.JSON(label="Embeddings")
 
-def greet(name):
-    return "Hello " + name + "!!"
+    # On submit, call embed() -> returns JSON
+    text_input.submit(embed, inputs=text_input, outputs=output)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+if __name__ == "__main__":
+    # queue() is optional but useful for concurrency
+    app.queue().launch(server_name="0.0.0.0", show_error=True, server_port=7860)