|
| 1 | +from typing import Any, Union |
| 2 | + |
| 3 | +import os |
| 4 | +import torch |
| 5 | +import numpy as np |
| 6 | +import pandas as pd |
| 7 | +from tqdm import tqdm |
| 8 | +from importlib.resources import files |
| 9 | +from .registry import register_evaluator, BaseEvaluator |
| 10 | +from .whowhat_metrics import EmbedsSimilarity |
| 11 | +from .utils import patch_awq_for_inference, get_ignore_parameters_flag |
| 12 | +from transformers import set_seed |
| 13 | +import datasets |
| 14 | +from torch import Tensor |
| 15 | + |
| 16 | +DEF_MAX_LENGTH = 100 |
| 17 | + |
| 18 | + |
| 19 | +def prepare_default_data(num_samples=None): |
| 20 | + DATASET_NAME = "microsoft/ms_marco" |
| 21 | + NUM_SAMPLES = num_samples if num_samples else 24 |
| 22 | + set_seed(42) |
| 23 | + default_dataset = datasets.load_dataset( |
| 24 | + DATASET_NAME, 'v2.1', split="test", streaming=True |
| 25 | + ).shuffle(42).take(NUM_SAMPLES) |
| 26 | + return default_dataset.map( |
| 27 | + lambda x: {'passages': x['passages']['passage_text']}, remove_columns=default_dataset.column_names |
| 28 | + ) |
| 29 | + |
| 30 | + |
| 31 | +def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor: |
| 32 | + left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0] |
| 33 | + if left_padding: |
| 34 | + return last_hidden_states[:, -1] |
| 35 | + else: |
| 36 | + sequence_lengths = attention_mask.sum(dim=1) - 1 |
| 37 | + batch_size = last_hidden_states.shape[0] |
| 38 | + batch_dim = torch.arange(batch_size, device=last_hidden_states.device) |
| 39 | + result = last_hidden_states[batch_dim, sequence_lengths] |
| 40 | + return result |
| 41 | + |
| 42 | + |
| 43 | +def mean_pooling(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor: |
| 44 | + input_mask_expanded = ( |
| 45 | + attention_mask.unsqueeze(-1).expand(last_hidden_states.size()).to(last_hidden_states.dtype) |
| 46 | + ) |
| 47 | + sum_embeddings = torch.sum(last_hidden_states * input_mask_expanded, 1) |
| 48 | + sum_mask = input_mask_expanded.sum(1) |
| 49 | + sum_mask = torch.clamp(sum_mask, min=1e-9) |
| 50 | + |
| 51 | + return sum_embeddings / sum_mask |
| 52 | + |
| 53 | + |
| 54 | +@register_evaluator( |
| 55 | + "text-embedding" |
| 56 | +) |
| 57 | +class EmbeddingsEvaluator(BaseEvaluator): |
| 58 | + def __init__( |
| 59 | + self, |
| 60 | + base_model: Any = None, |
| 61 | + tokenizer: Any = None, |
| 62 | + gt_data: str = None, |
| 63 | + test_data: Union[str, list] = None, |
| 64 | + num_samples=None, |
| 65 | + gen_embeds_fn=None, |
| 66 | + pooling_type=None, |
| 67 | + normalize=None, |
| 68 | + padding_side=None |
| 69 | + ) -> None: |
| 70 | + assert ( |
| 71 | + base_model is not None or gt_data is not None |
| 72 | + ), "Text generation pipeline for evaluation or ground trush data must be defined" |
| 73 | + |
| 74 | + self.test_data = test_data |
| 75 | + self.tokenizer = tokenizer |
| 76 | + self.num_samples = num_samples |
| 77 | + self.generation_fn = gen_embeds_fn |
| 78 | + self.pooling_type = pooling_type or 'cls' |
| 79 | + self.normalize = normalize or False |
| 80 | + self.padding_side = padding_side or 'right' |
| 81 | + self.gt_dir = os.path.dirname(gt_data) |
| 82 | + |
| 83 | + if base_model: |
| 84 | + self.gt_data = self._generate_data(base_model, gen_embeds_fn) |
| 85 | + else: |
| 86 | + self.gt_data = pd.read_csv(gt_data, keep_default_na=False) |
| 87 | + |
| 88 | + self.similarity = EmbedsSimilarity() |
| 89 | + self.last_cmp = None |
| 90 | + |
| 91 | + def get_generation_fn(self): |
| 92 | + return self.generation_fn |
| 93 | + |
| 94 | + def score(self, model_or_data, gen_answer_fn=None, output_dir=None, **kwargs): |
| 95 | + if output_dir is None: |
| 96 | + result_folder = os.path.join(self.gt_dir, "target") |
| 97 | + else: |
| 98 | + result_folder = os.path.join(output_dir, "target") |
| 99 | + |
| 100 | + if isinstance(model_or_data, str) and os.path.exists(model_or_data): |
| 101 | + predictions = pd.read_csv(model_or_data, keep_default_na=False) |
| 102 | + else: |
| 103 | + predictions = self._generate_data(model_or_data, gen_answer_fn, result_folder) |
| 104 | + self.predictions = predictions |
| 105 | + |
| 106 | + all_metrics_per_prompt = {} |
| 107 | + all_metrics = {} |
| 108 | + all_metrics, all_metrics_per_prompt = self.similarity.evaluate( |
| 109 | + self.gt_data, predictions |
| 110 | + ) |
| 111 | + |
| 112 | + self.last_cmp = all_metrics_per_prompt |
| 113 | + self.last_cmp["passages"] = predictions["passages"].values |
| 114 | + self.last_cmp["source_model"] = self.gt_data["embeds_path"].values |
| 115 | + self.last_cmp["optimized_model"] = predictions["embeds_path"].values |
| 116 | + self.last_cmp = pd.DataFrame(self.last_cmp) |
| 117 | + |
| 118 | + return pd.DataFrame(all_metrics_per_prompt), pd.DataFrame([all_metrics]) |
| 119 | + |
| 120 | + def worst_examples(self, top_k: int = 5, metric="similarity"): |
| 121 | + assert self.last_cmp is not None |
| 122 | + res = self.last_cmp.nsmallest(top_k, metric) |
| 123 | + return list(row for idx, row in res.iterrows()) |
| 124 | + |
| 125 | + def _generate_data(self, model, gen_answer_fn=None, result_dir="reference"): |
| 126 | + def default_gen_answer(model, tokenizer, passages, **kwargs): |
| 127 | + device = "cpu" |
| 128 | + if hasattr(model, "device"): |
| 129 | + device = model.device |
| 130 | + tokenizer_kwargs = {'padding': 'max_length', 'max_length': DEF_MAX_LENGTH, |
| 131 | + 'truncation': True, 'padding_side': kwargs.get('padding_side', 'right')} |
| 132 | + inputs = self.tokenizer(passages, return_tensors="pt", **tokenizer_kwargs).to(device) |
| 133 | + |
| 134 | + with torch.no_grad(): |
| 135 | + outputs = model(**inputs) |
| 136 | + |
| 137 | + if model.config.model_type == "qwen3" or kwargs.get("pooling_type", "last_token"): |
| 138 | + embeddings = last_token_pool(outputs.last_hidden_state, inputs["attention_mask"]) |
| 139 | + elif kwargs.get("pooling_type", "mean"): |
| 140 | + embeddings = mean_pooling(outputs.last_hidden_state, inputs["attention_mask"]) |
| 141 | + else: |
| 142 | + embeddings = outputs.last_hidden_state[:, 0] |
| 143 | + |
| 144 | + if kwargs.get("normalize", False): |
| 145 | + embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) |
| 146 | + return embeddings |
| 147 | + |
| 148 | + gen_answer_fn = gen_answer_fn or default_gen_answer |
| 149 | + |
| 150 | + if self.test_data: |
| 151 | + if isinstance(self.test_data, str): |
| 152 | + data = pd.read_csv(self.test_data) |
| 153 | + else: |
| 154 | + if isinstance(self.test_data, dict): |
| 155 | + assert "prompts" in self.test_data |
| 156 | + data = dict(self.test_data) |
| 157 | + else: |
| 158 | + data = {"prompts": list(self.test_data)} |
| 159 | + data = pd.DataFrame.from_dict(data) |
| 160 | + else: |
| 161 | + data = pd.DataFrame.from_dict(prepare_default_data(self.num_samples)) |
| 162 | + |
| 163 | + embeds_paths = [] |
| 164 | + passages = [] |
| 165 | + inptus = ( |
| 166 | + data.values |
| 167 | + if self.num_samples is None |
| 168 | + else data.values[: self.num_samples] |
| 169 | + ) |
| 170 | + |
| 171 | + if not os.path.exists(result_dir): |
| 172 | + os.makedirs(result_dir) |
| 173 | + |
| 174 | + for i, data in tqdm(enumerate(inptus), desc="Evaluate pipeline"): |
| 175 | + kwargs = {'padding_side': self.padding_side, |
| 176 | + 'pooling_type': self.pooling_type, |
| 177 | + 'normalize': self.normalize} |
| 178 | + result = gen_answer_fn(model, self.tokenizer, data[0], **kwargs) |
| 179 | + passages.append(data[0]) |
| 180 | + result_path = os.path.join(result_dir, f"embeds_{i}.npy") |
| 181 | + with open(result_path, 'wb') as f: |
| 182 | + np.save(f, result) |
| 183 | + embeds_paths.append(result_path) |
| 184 | + |
| 185 | + res_data = {"passages": passages, "embeds_path": embeds_paths} |
| 186 | + df = pd.DataFrame(res_data) |
| 187 | + |
| 188 | + return df |
0 commit comments