Spaces:
Running
Running
File size: 1,118 Bytes
b82a487 77a97ce b82a487 77a97ce 0d77bb1 77a97ce 0d77bb1 77a97ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
from enum import Enum
import safetensors
import safetensors.torch
import torch
import wandb
class SimilarityMetric(Enum):
COSINE = "cosine"
EUCLIDEAN = "euclidean"
def mean_pooling(token_embeddings, mask):
token_embeddings = token_embeddings.masked_fill(~mask[..., None].bool(), 0.0)
sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None]
return sentence_embeddings
def argsort_scores(scores: list[float], descending: bool = False):
return [
{"item": item, "original_index": idx}
for idx, item in sorted(
list(enumerate(scores)), key=lambda x: x[1], reverse=descending
)
]
def save_vector_index(
vector_index: torch.Tensor,
type: str,
index_name: str,
metadata: dict,
filename: str = "vector_index.safetensors",
):
safetensors.torch.save_file({"vector_index": vector_index.cpu()}, filename)
if wandb.run:
artifact = wandb.Artifact(
name=index_name,
type=type,
metadata=metadata,
)
artifact.add_file(filename)
artifact.save()
|