chore: import upstream snapshot with attribution

2026-07-13 13:24:13 +08:00
commit 1037506f2e
6050 changed files with 1731598 additions and 0 deletions
@@ -0,0 +1,119 @@
+import os
+import tqdm
+import torch
+
+from contextlib import nullcontext
+from torch.utils.data import DataLoader
+from functools import partial
+from datasets import load_dataset
+from typing import Dict, List
+from transformers.file_utils import PaddingStrategy
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    HfArgumentParser,
+    BatchEncoding
+)
+
+from config import Arguments
+from logger_config import logger
+from utils import move_to_cuda
+from models import BiencoderModelForInference, BiencoderOutput
+
+parser = HfArgumentParser((Arguments,))
+args: Arguments = parser.parse_args_into_dataclasses()[0]
+
+
+def _psg_transform_func(tokenizer: PreTrainedTokenizerFast,
+                        examples: Dict[str, List]) -> BatchEncoding:
+    batch_dict = tokenizer(examples['title'],
+                           text_pair=examples['contents'],
+                           max_length=args.p_max_len,
+                           padding=PaddingStrategy.DO_NOT_PAD,
+                           truncation=True)
+    # for co-Condenser reproduction purpose only
+    if args.model_name_or_path.startswith('Luyu/'):
+        del batch_dict['token_type_ids']
+
+    return batch_dict
+
+
+@torch.no_grad()
+def _worker_encode_passages(gpu_idx: int):
+    def _get_out_path(shard_idx: int = 0) -> str:
+        return '{}/shard_{}_{}'.format(args.encode_save_dir, gpu_idx, shard_idx)
+
+    if os.path.exists(_get_out_path(0)):
+        logger.error('{} already exists, will skip encoding'.format(_get_out_path(0)))
+        return
+
+    dataset = load_dataset('json', data_files=args.encode_in_path)['train']
+    if args.dry_run:
+        dataset = dataset.select(range(4096))
+    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
+                            index=gpu_idx,
+                            contiguous=True)
+
+    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
+    torch.cuda.set_device(gpu_idx)
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: BiencoderModelForInference = BiencoderModelForInference.build(args)
+    model.eval()
+    model.cuda()
+
+    dataset.set_transform(partial(_psg_transform_func, tokenizer))
+
+    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if args.fp16 else None)
+    data_loader = DataLoader(
+        dataset,
+        batch_size=args.encode_batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=data_collator,
+        pin_memory=True)
+
+    num_encoded_docs, encoded_embeds, cur_shard_idx = 0, [], 0
+    for batch_dict in tqdm.tqdm(data_loader, desc='passage encoding', mininterval=8):
+        batch_dict = move_to_cuda(batch_dict)
+
+        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
+            outputs: BiencoderOutput = model(query=None, passage=batch_dict)
+        encoded_embeds.append(outputs.p_reps.cpu())
+        num_encoded_docs += outputs.p_reps.shape[0]
+
+        if num_encoded_docs >= args.encode_shard_size:
+            out_path = _get_out_path(cur_shard_idx)
+            concat_embeds = torch.cat(encoded_embeds, dim=0)
+            logger.info('GPU {} save {} embeds to {}'.format(gpu_idx, concat_embeds.shape[0], out_path))
+            torch.save(concat_embeds, out_path)
+
+            cur_shard_idx += 1
+            num_encoded_docs = 0
+            encoded_embeds.clear()
+
+    if num_encoded_docs > 0:
+        out_path = _get_out_path(cur_shard_idx)
+        concat_embeds = torch.cat(encoded_embeds, dim=0)
+        logger.info('GPU {} save {} embeds to {}'.format(gpu_idx, concat_embeds.shape[0], out_path))
+        torch.save(concat_embeds, out_path)
+
+    logger.info('Done computing score for worker {}'.format(gpu_idx))
+
+
+def _batch_encode_passages():
+    logger.info('Args={}'.format(str(args)))
+    gpu_count = torch.cuda.device_count()
+    if gpu_count == 0:
+        logger.error('No gpu available')
+        return
+
+    logger.info('Use {} gpus'.format(gpu_count))
+    torch.multiprocessing.spawn(_worker_encode_passages, args=(), nprocs=gpu_count)
+    logger.info('Done batch encode passages')
+
+
+if __name__ == '__main__':
+    _batch_encode_passages()