Merge pull request #642 from 545999961/master

update reranker v2

Author: 545999961

FlagEmbedding/flag_reranker.py

@@ -4,7 +4,7 @@
 import torch
 from torch import Tensor
 from torch.utils.data import DataLoader
-from tqdm import tqdm
+from tqdm import tqdm, trange
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, is_torch_npu_available
 
 import warnings
@@ -269,32 +269,96 @@ def __init__(
 
     @torch.no_grad()
     def compute_score(self, sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]], batch_size: int = 16,
-                      max_length: int = 512, prompt: str = None, normalize: bool = False) -> List[float]:
+                      max_length: int = 512, prompt: str = None, normalize: bool = False,
+                      use_dataloader: bool = True, num_workers: int = None) -> List[float]:
         assert isinstance(sentence_pairs, list)
         if isinstance(sentence_pairs[0], str):
             sentence_pairs = [sentence_pairs]
 
         length_sorted_idx = np.argsort([-self._text_length(q) - self._text_length(p) for q, p in sentence_pairs])
         sentences_sorted = [sentence_pairs[idx] for idx in length_sorted_idx]
 
-        dataset = DatasetForReranker(sentences_sorted,
-                                     self.model_name_or_path,
-                                     max_length,
-                                     cache_dir=self.cache_dir,
-                                     prompt=prompt)
-        dataloader = DataLoader(dataset, shuffle=False, batch_size=batch_size, drop_last=False,
-                                num_workers=min(batch_size, 16),
-                                collate_fn=collater(self.tokenizer, max_length))
+        dataset, dataloader = None, None
+        if use_dataloader:
+            if num_workers is None:
+                num_workers = min(batch_size, 16)
+                dataset = DatasetForReranker(sentences_sorted,
+                                             self.model_name_or_path,
+                                             max_length,
+                                             cache_dir=self.cache_dir,
+                                             prompt=prompt)
+                dataloader = DataLoader(dataset, shuffle=False, batch_size=batch_size, drop_last=False,
+                                        num_workers=num_workers,
+                                        collate_fn=collater(self.tokenizer, max_length))
 
         all_scores = []
-        for inputs in tqdm(dataloader):
-            inputs = inputs.to(self.device)
-
-            outputs = self.model(**inputs, output_hidden_states=True)
-            logits = outputs.logits
-            scores = last_logit_pool(logits, inputs['attention_mask'])
-            scores = scores[:, self.yes_loc]
-            all_scores.extend(scores.cpu().float().tolist())
+        if dataloader is not None:
+            for inputs in tqdm(dataloader):
+                inputs = inputs.to(self.device)
+
+                outputs = self.model(**inputs, output_hidden_states=True)
+                logits = outputs.logits
+                scores = last_logit_pool(logits, inputs['attention_mask'])
+                scores = scores[:, self.yes_loc]
+                all_scores.extend(scores.cpu().float().tolist())
+        else:
+            if prompt is None:
+                prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
+            prompt_inputs = self.tokenizer(prompt,
+                                                return_tensors=None,
+                                                add_special_tokens=False)['input_ids']
+            sep = "\n"
+            sep_inputs = self.tokenizer(sep,
+                                             return_tensors=None,
+                                             add_special_tokens=False)['input_ids']
+            encode_max_length = max_length + len(sep_inputs) + len(prompt_inputs)
+            for batch_start in trange(0, len(sentences_sorted), batch_size):
+                batch_sentences = sentences_sorted[batch_start:batch_start + batch_size]
+                batch_sentences = [(f'A: {q}', f'B: {p}') for q,p in batch_sentences]
+                queries = [s[0] for s in batch_sentences]
+                passages = [s[1] for s in batch_sentences]
+                queries_inputs = self.tokenizer(queries,
+                                                return_tensors=None,
+                                                add_special_tokens=False,
+                                                max_length=max_length * 3 // 4,
+                                                truncation=True)
+                passages_inputs = self.tokenizer(passages,
+                                                 return_tensors=None,
+                                                 add_special_tokens=False,
+                                                 max_length=max_length,
+                                                 truncation=True)
+
+                batch_inputs = []
+                for query_inputs, passage_inputs in zip(queries_inputs['input_ids'], passages_inputs['input_ids']):
+                    item = self.tokenizer.prepare_for_model(
+                        [self.tokenizer.bos_token_id] + query_inputs,
+                        sep_inputs + passage_inputs,
+                        truncation='only_second',
+                        max_length=encode_max_length,
+                        padding=False,
+                        return_attention_mask=False,
+                        return_token_type_ids=False,
+                        add_special_tokens=False
+                    )
+                    item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
+                    item['attention_mask'] = [1] * len(item['input_ids'])
+                    item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
+                    if 'position_ids' in item.keys():
+                        item['position_ids'] = list(range(len(item['input_ids'])))
+                    batch_inputs.append(item)
+
+                collater_instance = collater(self.tokenizer, max_length)
+                batch_inputs = collater_instance(
+                    [{'input_ids': item['input_ids'], 'attention_mask': item['attention_mask']} for item in
+                     batch_inputs])
+
+                batch_inputs = {key: val.to(self.device) for key, val in batch_inputs.items()}
+
+                outputs = self.model(**batch_inputs, output_hidden_states=True)
+                logits = outputs.logits
+                scores = last_logit_pool(logits, batch_inputs['attention_mask'])
+                scores = scores[:, self.yes_loc]
+                all_scores.extend(scores.cpu().float().tolist())
 
         all_scores = [all_scores[idx] for idx in np.argsort(length_sorted_idx)]
 
@@ -323,6 +387,7 @@ def _text_length(self, text: Union[List[int], List[List[int]]]):
         else:
             return sum([len(t) for t in text])  # Sum of length of individual strings
 
+
 class LayerWiseFlagLLMReranker:
     def __init__(
             self,
@@ -378,40 +443,112 @@ def __init__(
     @torch.no_grad()
     def compute_score(self, sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]], batch_size: int = 16,
                       max_length: int = 512, cutoff_layers: List[int] = None, prompt: str = None,
-                      normalize: bool = False) -> Union[float, List[float], List[List[float]]]:
+                      normalize: bool = False, use_dataloader: bool = True,
+                      num_workers: int = None) -> Union[float, List[float], List[List[float]]]:
         assert isinstance(sentence_pairs, list)
         if isinstance(sentence_pairs[0], str):
             sentence_pairs = [sentence_pairs]
 
         length_sorted_idx = np.argsort([-self._text_length(q) - self._text_length(p) for q, p in sentence_pairs])
         sentences_sorted = [sentence_pairs[idx] for idx in length_sorted_idx]
 
-        dataset = DatasetForReranker(sentences_sorted,
-                                     self.model_name_or_path,
-                                     max_length,
-                                     cache_dir=self.cache_dir,
-                                     prompt=prompt)
-        dataloader = DataLoader(dataset, shuffle=False, batch_size=batch_size, drop_last=False,
-                                num_workers=min(batch_size, 16),
-                                collate_fn=collater(self.tokenizer, max_length))
+        dataset, dataloader = None, None
+        if use_dataloader:
+            if num_workers is None:
+                num_workers = min(batch_size, 16)
+                dataset = DatasetForReranker(sentences_sorted,
+                                             self.model_name_or_path,
+                                             max_length,
+                                             cache_dir=self.cache_dir,
+                                             prompt=prompt)
+                dataloader = DataLoader(dataset, shuffle=False, batch_size=batch_size, drop_last=False,
+                                        num_workers=num_workers,
+                                        collate_fn=collater(self.tokenizer, max_length))
 
         all_scores = []
-        for inputs in tqdm(dataloader):
-            inputs = inputs.to(self.device)
-
-            outputs = self.model(**inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
-            all_logits = outputs.logits
-            tmp_all_scores = []
-            for logits in all_logits:
-                scores = last_logit_pool_layerwise(logits, inputs['attention_mask'])
-                tmp_all_scores.append(scores.contiguous())
-
-            if len(all_scores) == 0:
-                for _ in range(len(tmp_all_scores)):
-                    all_scores.append([])
-
-            for i in range(len(tmp_all_scores)):
-                all_scores[i].extend(tmp_all_scores[i].cpu().float().tolist())
+        if dataloader is not None:
+            for inputs in tqdm(dataloader):
+                inputs = inputs.to(self.device)
+
+                outputs = self.model(**inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
+                all_logits = outputs.logits
+                tmp_all_scores = []
+                for logits in all_logits:
+                    scores = last_logit_pool_layerwise(logits, inputs['attention_mask'])
+                    tmp_all_scores.append(scores.contiguous())
+
+                if len(all_scores) == 0:
+                    for _ in range(len(tmp_all_scores)):
+                        all_scores.append([])
+
+                for i in range(len(tmp_all_scores)):
+                    all_scores[i].extend(tmp_all_scores[i].cpu().float().tolist())
+        else:
+            if prompt is None:
+                prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
+            prompt_inputs = self.tokenizer(prompt,
+                                                return_tensors=None,
+                                                add_special_tokens=False)['input_ids']
+            sep = "\n"
+            sep_inputs = self.tokenizer(sep,
+                                             return_tensors=None,
+                                             add_special_tokens=False)['input_ids']
+            encode_max_length = max_length + len(sep_inputs) + len(prompt_inputs)
+            for batch_start in trange(0, len(sentences_sorted), batch_size):
+                batch_sentences = sentences_sorted[batch_start:batch_start + batch_size]
+                batch_sentences = [(f'A: {q}', f'B: {p}') for q, p in batch_sentences]
+                queries = [s[0] for s in batch_sentences]
+                passages = [s[1] for s in batch_sentences]
+                queries_inputs = self.tokenizer(queries,
+                                                return_tensors=None,
+                                                add_special_tokens=False,
+                                                max_length=max_length * 3 // 4,
+                                                truncation=True)
+                passages_inputs = self.tokenizer(passages,
+                                                 return_tensors=None,
+                                                 add_special_tokens=False,
+                                                 max_length=max_length,
+                                                 truncation=True)
+
+                batch_inputs = []
+                for query_inputs, passage_inputs in zip(queries_inputs['input_ids'], passages_inputs['input_ids']):
+                    item = self.tokenizer.prepare_for_model(
+                        [self.tokenizer.bos_token_id] + query_inputs,
+                        sep_inputs + passage_inputs,
+                        truncation='only_second',
+                        max_length=encode_max_length,
+                        padding=False,
+                        return_attention_mask=False,
+                        return_token_type_ids=False,
+                        add_special_tokens=False
+                    )
+                    item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
+                    item['attention_mask'] = [1] * len(item['input_ids'])
+                    item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
+                    if 'position_ids' in item.keys():
+                        item['position_ids'] = list(range(len(item['input_ids'])))
+                    batch_inputs.append(item)
+
+                collater_instance = collater(self.tokenizer, max_length)
+                batch_inputs = collater_instance(
+                    [{'input_ids': item['input_ids'], 'attention_mask': item['attention_mask']} for item in
+                     batch_inputs])
+
+                batch_inputs = {key: val.to(self.device) for key, val in batch_inputs.items()}
+
+                outputs = self.model(**batch_inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
+                all_logits = outputs.logits
+                tmp_all_scores = []
+                for logits in all_logits:
+                    scores = last_logit_pool_layerwise(logits, batch_inputs['attention_mask'])
+                    tmp_all_scores.append(scores.contiguous())
+
+                if len(all_scores) == 0:
+                    for _ in range(len(tmp_all_scores)):
+                        all_scores.append([])
+
+                for i in range(len(tmp_all_scores)):
+                    all_scores[i].extend(tmp_all_scores[i].cpu().float().tolist())
 
         for i in range(len(all_scores)):
             all_scores[i] = [all_scores[i][idx] for idx in np.argsort(length_sorted_idx)]

FlagEmbedding/llm_reranker/finetune_for_instruction/data.py

@@ -1,5 +1,6 @@
 import re
 import sys
+from typing import List
 
 import math
 import os.path
@@ -60,7 +61,7 @@ def is_chinese(self, text):
         chinese_pattern = re.compile('[\u4e00-\u9fa5]')
         return bool(chinese_pattern.search(text))
 
-    def __getitem__(self, item) -> list[BatchEncoding]:
+    def __getitem__(self, item) -> List[BatchEncoding]:
         query = self.dataset[item]['query']
 
         passages = []

FlagEmbedding/llm_reranker/finetune_for_layerwise/arguments.py

@@ -71,7 +71,7 @@ class ModelArguments:
     from_peft: str = field(
         default=None
     )
-    lora_extra_parameters: str = field(
+    lora_extra_parameters: Optional[List[str]] = field(
         default=None
     )
     start_layer: int = field(

FlagEmbedding/llm_reranker/finetune_for_layerwise/data.py

@@ -1,5 +1,6 @@
 import re
 import sys
+from typing import List
 
 import math
 import os.path
@@ -56,7 +57,7 @@ def __init__(
     def __len__(self):
         return self.total_len
 
-    def __getitem__(self, item) -> list[BatchEncoding]:
+    def __getitem__(self, item) -> List[BatchEncoding]:
         query = self.dataset[item]['query']
 
         passages = []