from rank_bm25 import BM25Okapi import os import json import re from transformers import AutoTokenizer, AutoModel import argparse from tqdm import tqdm import torch import argparse # ========================= # 配置多卡 / 设备 # ========================= device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_path = "./bge-base-en-v1.5" # 或 "BAAI/bge-base-en-v1.5" tokenizer = AutoTokenizer.from_pretrained(model_path) # ========================= # 解析命令行参数 # ========================= parser = argparse.ArgumentParser() parser.add_argument( "--task", type=int, default=None, help="Process a single task (1–7 except 6)" ) parser.add_argument( "--tasks", type=int, nargs="*", default=None, help="Process multiple tasks, e.g., --tasks 1 3 5" ) parser.add_argument( "--data_dir", type=str, required=True, help="Base directory for LaMP time datasets, e.g. /opt/tiger/PLLM2.0/dataset" ) args = parser.parse_args() # 决定最终 task 列表 if args.task is not None: task_list = [args.task] elif args.tasks is not None: task_list = args.tasks else: # 默认仍然处理所有任务（保持原逻辑） task_list = [1,2,3,4,5,7] print(f"📌 Tasks to run: {task_list}") # 只在有 GPU 时用 bf16，没有 GPU 保持默认 float32，避免行为变化 if torch.cuda.is_available(): model = AutoModel.from_pretrained( model_path, torch_dtype=torch.bfloat16, ) else: model = AutoModel.from_pretrained(model_path) if torch.cuda.is_available(): num_gpus = torch.cuda.device_count() print(f"✅ Detected {num_gpus} GPUs") if num_gpus > 1: print(f"✅ Using DataParallel on {num_gpus} GPUs") model = torch.nn.DataParallel(model) # 自动多卡并行 model.to(device) else: print("⚠️ CUDA not available, using CPU") model.to(device) # 推理模式（不会改逻辑，只关掉 dropout 等） model.eval() def process(idx, entry): question = entry["input"] profile = entry["profile"] new_entry = {} if idx == 1: substr = 'title "' plc = question.find(substr) + len(substr) substr2 = '",' plc2 = question.find(substr2, plc) real_question = question[plc:plc2] new_entry["question"] = real_question new_entry["profile_txt"] = [] new_entry["profile_id"] = [] for his in profile: new_entry["profile_txt"].append(his["title"] + " " + his["abstract"]) new_entry["profile_id"].append(his["id"]) if idx == 2: substr = "description: " plc = question.find(substr) + len(substr) real_question = question[plc:] new_entry["question"] = real_question new_entry["profile_txt"] = [] new_entry["profile_id"] = [] for his in profile: new_entry["profile_txt"].append(his["description"]) new_entry["profile_id"].append(his["id"]) if idx == 3: substr = "review: " plc = question.find(substr) + len(substr) real_question = question[plc:] new_entry["question"] = real_question new_entry["profile_txt"] = [] new_entry["profile_id"] = [] for his in profile: new_entry["profile_txt"].append(his["text"]) new_entry["profile_id"].append(his["id"]) if idx == 4: substr = "article: " plc = question.find(substr) + len(substr) real_question = question[plc:] new_entry["question"] = real_question new_entry["profile_txt"] = [] new_entry["profile_id"] = [] for his in profile: new_entry["profile_txt"].append(his["text"] + " " + his["title"]) new_entry["profile_id"].append(his["id"]) if idx == 5: substr = "paper: " plc = question.find(substr) + len(substr) real_question = question[plc:] new_entry["question"] = real_question new_entry["profile_txt"] = [] new_entry["profile_id"] = [] for his in profile: new_entry["profile_txt"].append(his["title"] + " " + his["abstract"]) new_entry["profile_id"].append(his["id"]) if idx == 7: substr = "before or after it: " plc = question.find(substr) + len(substr) real_question = question[plc:] new_entry["question"] = real_question new_entry["profile_txt"] = [] new_entry["profile_id"] = [] for his in profile: new_entry["profile_txt"].append(his["text"]) new_entry["profile_id"].append(his["id"]) return new_entry remove_chars = r'[·’!"\#$%&\'()＃！（）*+,-./:;<=>?\@，：?￥★、…．＞【】［］《》？“”‘’\[\\]^_`{|}~]+' def sort(corpus_embedding, query_embedding, his_len, new_entry): all_list = [] for i in range(his_len): all_list.append( ( float(corpus_embedding[i] @ query_embedding), new_entry["profile_txt"][i], new_entry["profile_id"][i], ) ) all_list.sort(key=lambda x: -x[0]) return all_list @torch.no_grad() def get_embedding(sentences, batch_size=4096): print(f"Total sentences: {len(sentences)}") 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 embeddings_list = [] for i in tqdm(range(0, len(sentences), batch_size)): batch_sentences = sentences[i : i + batch_size] inputs = tokenizer( batch_sentences, padding=True, truncation=True, return_tensors="pt", ).to(device) # ✅ 一次性搬到 device，多卡自动拆 outputs = model(**inputs) if isinstance(outputs, (list, tuple)): token_embeddings = outputs[0] else: token_embeddings = outputs.last_hidden_state # 逻辑保持不变：仍然用 CLS 向量 batch_emb = torch.nn.functional.normalize( token_embeddings[:, 0], p=2, dim=1 ).detach().cpu() embeddings_list.append(batch_emb) # 只删引用就够了，不再 empty_cache del inputs, outputs, token_embeddings, batch_emb # ❌ 不再调用 torch.cuda.empty_cache() embeddings = torch.cat(embeddings_list, dim=0) return embeddings os.makedirs("./bge_emb", exist_ok=True) for idx in task_list: if idx == 6: print("⚠️ Task 6 does not exist, skipping.") continue dir_name = os.path.join(args.data_dir, f"LaMP_time_{idx}") # ========== train ========== file_name = os.path.join(dir_name, "train_questions.json") print(file_name) dataset = json.load(open(file_name)) all_txt_list = [] for entry in dataset: new_entry = process(idx, entry) all_txt_list.extend(new_entry["profile_txt"] + [new_entry["question"]]) corpus_embedding_all_list = get_embedding(all_txt_list) torch.save(corpus_embedding_all_list, f"./bge_emb/task_{idx}_train_bge.emb") # ========== dev ========== file_name = os.path.join(dir_name, "dev_questions.json") print(file_name) dataset = json.load(open(file_name)) all_txt_list = [] for entry in dataset: new_entry = process(idx, entry) all_txt_list.extend(new_entry["profile_txt"] + [new_entry["question"]]) corpus_embedding_all_list = get_embedding(all_txt_list) torch.save(corpus_embedding_all_list, f"./bge_emb/task_{idx}_dev_bge.emb")