[auto] Upload file Recommend/embedding.py

88ca4c3 verified 11 months ago

7.79 kB

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")

	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")