FlagEmbedding模型微调时json数据处理

from sentence_transformers import CrossEncoder from sentence_transformers import SentenceTransformer from sentence_transformers.util import cos_sim from transformers import AutoModelForSequenceClassification, AutoTokenizer, BatchEncoding, PreTrainedTokenizerFast import torch import json import numpy as np def is_within_range(element): col = element[0][0] #print("col", col) if col > 5: return False return True def read_data(file_path): data_list = [] with open(file_path, 'r', encoding='utf-8') as file: for lines in file: a = [] data = json.loads(lines) temp = [] temp.append(data['query']) temp.append(data['pos'][0]) #print('temp is:', temp, '\n') a.append(temp) for i in range(0, len(data['neg'])): temp2 = [] temp2.append(data['query']) temp2.append(data['neg'][i]) #print('temp2 is:', temp2, '\n') a.append(temp2) data_list.append(a) # print(a, '\n') # print(data_list, '\n') return data_list def base_large_v2m3(model_path, sentences): # 获取相关性分数（分数越高表明相关性越高） from FlagEmbedding import FlagReranker from datetime import datetime now = datetime.now() reranker = FlagReranker(model_path, use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation end = datetime.now() print("加载模型耗时：", (end - now).total_seconds()) num = len(sentences) # scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]) # print("bge reranker scores is:", scores) # [-8.1875, 5.26171875] # You can map the scores into 0-1 by set "normalize=True", which will apply sigmoid function to the score now2 = datetime.now() scores = reranker.compute_score(sentences, normalize=True) end2 = datetime.now() avg_time = (end2 - now2).total_seconds()/num #print("计算每一条样例中，每一对向量的相关性耗时：", (end2 - now2).total_seconds()/num) #print("mormalized bge reranker scores is:", scores) # [0.00027803096387751553, 0.9948403768236574] return scores, avg_time def llm_reranker(model_path, sentences): from FlagEmbedding import FlagLLMReranker from datetime import datetime now = datetime.now() reranker = FlagLLMReranker(model_path, use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation # reranker = FlagLLMReranker('BAAI/bge-reranker-v2-gemma', use_bf16=True) # You can also set use_bf16=True to speed up computation with a slight performance degradation end = datetime.now() print("加载模型耗时：", (end - now).total_seconds()) num = len(sentences) now2 = datetime.now() # score = reranker.compute_score(sentences) # print("llm reranker scores is:", score) scores = reranker.compute_score(sentences, batch_size=8) end2 = datetime.now() avg_time = (end2 - now2).total_seconds()/num #print("normalized llm reranker scores is:", scores) return scores, avg_time def layer_reranker(model_path, sentences): from FlagEmbedding import LayerWiseFlagLLMReranker from datetime import datetime now = datetime.now() reranker = LayerWiseFlagLLMReranker(model_path, use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation end = datetime.now() print("加载模型耗时：", (end - now).total_seconds()) num = len(sentences) # reranker = LayerWiseFlagLLMReranker('BAAI/bge-reranker-v2-minicpm-layerwise', use_bf16=True) # You can also set use_bf16=True to speed up computation with a slight performance degradation #score = reranker.compute_score(sentences, cutoff_layers=[28]) # Adjusting 'cutoff_layers' to pick which layers are used for computing the score. #print("layer reranker scores is:", score) now2 = datetime.now() scores = reranker.compute_score(sentences, cutoff_layers=[28]) end2 = datetime.now() avg_time = (end2 - now2).total_seconds()/num #print("normalized llm reranker scores is:", scores) return scores, avg_time def hf_reranker(model_path, sentences): import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(model_path) model = AutoModelForSequenceClassification.from_pretrained(model_path) model.eval() with torch.no_grad(): inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt', max_length=512) scores = model(**inputs, return_dict=True).logits.view(-1, ).float() print("hf load reranker scores is:", scores) def llm_reranker2(model_path, sentences): import torch from transformers import AutoModelForCausalLM, AutoTokenizer def get_inputs(sentences, tokenizer, prompt=None, max_length=1024): 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'." sep = "\n" prompt_inputs = tokenizer(prompt, return_tensors=None, add_special_tokens=False)['input_ids'] sep_inputs = tokenizer(sep, return_tensors=None, add_special_tokens=False)['input_ids'] inputs = [] for query, passage in sentences: query_inputs = tokenizer(f'A: {query}', return_tensors=None, add_special_tokens=False, max_length=max_length * 3 // 4, truncation=True) passage_inputs = tokenizer(f'B: {passage}',return_tensors=None,add_special_tokens=False,max_length=max_length,truncation=True) item = tokenizer.prepare_for_model( [tokenizer.bos_token_id] + query_inputs['input_ids'], sep_inputs + passage_inputs['input_ids'], truncation='only_second',max_length=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']) inputs.append(item) return tokenizer.pad(inputs,padding=True, max_length=max_length + len(sep_inputs) + len(prompt_inputs), pad_to_multiple_of=8,return_tensors='pt',) tokenizer = AutoTokenizer.from_pretrained(model_path) model = AutoModelForCausalLM.from_pretrained(model_path) yes_loc = tokenizer('Yes', add_special_tokens=False)['input_ids'][0] model.eval() with torch.no_grad(): inputs = get_inputs(sentences, tokenizer) scores = model(**inputs, return_dict=True).logits[:, -1, yes_loc].view(-1, ).float() print("llm reranker scores2 is:", scores) def layer_reranker2(model_path, sentences): import torch from transformers import AutoModelForCausalLM, AutoTokenizer def get_inputs(sentences, tokenizer, prompt=None, max_length=1024): 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'." sep = "\n" prompt_inputs = tokenizer(prompt,return_tensors=None,add_special_tokens=False)['input_ids'] sep_inputs = tokenizer(sep,return_tensors=None,add_special_tokens=False)['input_ids'] inputs = [] for query, passage in sentences: query_inputs = tokenizer(f'A: {query}',return_tensors=None,add_special_tokens=False,max_length=max_length * 3 // 4, truncation=True) passage_inputs = tokenizer(f'B: {passage}'