Create utils.py

utils.py

@@ -0,0 +1,151 @@
+import numpy as np
+import torch
+
+from typing import Any
+from transformers import AutoTokenizer
+
+
+def splade_max(features, attention_mask):
+    """
+    SPLADE pooling operation
+    """
+    relu = torch.nn.ReLU(inplace=False)
+    values, ids_ = torch.max(
+        torch.log(1 + relu(features)) * attention_mask.unsqueeze(-1), dim=1
+    )
+    return values, ids_
+
+
+def encode(
+    self,
+    sentences: list[str],
+    max_length: int = 1024,
+    prompt_type: str = "document",
+    return_dict: bool = False,
+    print_dict: bool = False,
+    batch_size: int = 8,
+    top_k_q: int = -1,
+    top_k_d: int = -1,
+    **kwargs: Any,
+) -> np.ndarray:
+    all_embeddings = []
+    for i in range(0, len(sentences), batch_size):
+        batch_texts = sentences[i : i + batch_size]
+        batch_dict = self.create_batch_dict(batch_texts, max_length)
+        batch_dict = {
+            key: value.to(self.model.device) for key, value in batch_dict.items()
+        }
+        with torch.no_grad():
+            splare_reps = self(**batch_dict)[0]
+            if prompt_type == "query" and top_k_q > 0:
+                splare_reps = top_k(splare_reps, top_k_q)
+            if prompt_type == "document" and top_k_d > 0:
+                splare_reps = top_k(splare_reps, top_k_d)
+            all_embeddings.append(splare_reps.cpu().float().numpy())
+    if return_dict:
+        d = bow_dict(self, np.concatenate(all_embeddings, axis=0))
+        if print_dict:
+            print_bow_bars(sentences, d)
+        return d
+    else:
+        return np.concatenate(all_embeddings, axis=0)
+
+
+def bow_dict(self, embeddings):
+    out = []
+    for vector in embeddings:
+        idx = np.nonzero(vector)[0]
+        weights = vector[idx]
+        d = {k: v for k, v in zip(idx.tolist(), weights.tolist())}
+        sorted_d = {
+            self.reverse_voc[k]: float(v)
+            for k, v in sorted(d.items(), key=lambda item: item[1], reverse=True)
+        }
+        out.append(sorted_d)
+    return out
+
+
+def print_bow_bars(sentences, bow_list, width=20):
+    ascii_header("TOP ACTIVATED WORDS")
+    for sent, bow in zip(sentences, bow_list):
+        print(f"* INPUT: {sent}\n")
+        max_w = max(bow.values())
+        for k, v in sorted(bow.items(), key=lambda x: x[1], reverse=True):
+            bar = "█" * int(v / max_w * width)
+            print(f"{k[:25]:25} | {bar} {v:.2f}")
+        print("\n")
+
+
+def ascii_header(title, width=70):
+    title = f" {title} "
+    print("+" + "-" * (width - 2) + "+")
+    print("|" + title.center(width - 2) + "|")
+    print("+" + "-" * (width - 2) + "+")
+    print("\n")
+
+
+def similarity(self, a, b) -> torch.Tensor:
+    """
+    MTEB eval requires this
+    """
+    if not isinstance(a, torch.Tensor):
+        a = torch.tensor(a)
+    if not isinstance(b, torch.Tensor):
+        b = torch.tensor(b)
+
+    def _dot_score_core(a_tensor, b_tensor):
+        if len(a_tensor.shape) == 1:
+            a_tensor = a_tensor.unsqueeze(0)
+        if len(b_tensor.shape) == 1:
+            b_tensor = b_tensor.unsqueeze(0)
+        return a_tensor @ b_tensor.transpose(0, 1)
+
+    return _dot_score_core(a, b)
+
+
+def prepare_tokenizer(tokenizer_name: str, padding_side="right"):
+    """
+    loads and prepares tokenizer
+    """
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    tokenizer.pad_token = (
+        tokenizer.bos_token or tokenizer.pad_token or tokenizer.eos_token
+    )
+    tokenizer.padding_side = padding_side
+    return tokenizer
+
+
+def get_decoder_model(
+    model_name_or_path: str, attn_implementation: str, bidirectional: bool, base_cfg
+):
+    """
+    base_cfg is the pretrained config of the underlying model
+    """
+    print("WARNING: bidirectional only tested for transformer 4.51.2")
+    assert (
+        bidirectional is True
+    ), "the model has been trained with bi-directional attention!"
+    assert (
+        attn_implementation == "flash_attention_2"
+    ), f"bidir models only support flash_attention_2 for now, not {attn_implementation}!"
+    from .modeling_qwen3_bidir import Qwen3BidirForCausalLM
+
+    return Qwen3BidirForCausalLM.from_pretrained(
+        model_name_or_path,
+        config=base_cfg,
+        torch_dtype=torch.bfloat16,
+        attn_implementation=attn_implementation,
+    )
+
+
+def top_k(x: torch.Tensor, k: int) -> torch.Tensor:
+    """
+    zeroes out all but the top-k values in the last dimension of x
+    """
+    _, topk_indices = x.topk(k, dim=-1)
+    # create a zero tensor of the same shape as x
+    mask = torch.zeros_like(x, dtype=torch.bool)
+    # use scatter along the last dimension
+    mask.scatter_(-1, topk_indices, True)
+    # zero out all but the top-k
+    return x * mask