from transformers import pipeline import torch import pandas as pd import sys import os sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'document_retrieval', 'Decompose_retrieval')) import ragqa_paths # [ragqa] portable paths from vllm import LLM, SamplingParams from transformers import AutoTokenizer from openai import OpenAI from qa_webq import * import requests import pickle from tqdm import tqdm import os os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' os.environ['https_proxy'] = 'http://127.0.0.1:7890/' os.environ['http_proxy'] = 'http://127.0.0.1:7890/' dataset = 'triviaqa' def call_llama3_single_prompt( inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=20, temperature=0.0 ): inputs_ls = [] if isinstance(inputs, str): messages = [ {"role": "user", "content": inputs}, ] inputs_ls.append(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)) else: for idx in range(len(inputs)): inputs_ls.append(tokenizer.apply_chat_template(inputs[idx], tokenize=False, add_generation_prompt=True)) # ans = get_vllm_llama(temperature, max_decode_steps, inputs_ls) results = client.completions.create( model= ragqa_paths.LLAMA_MODEL, max_tokens=max_decode_steps, temperature=0, prompt=inputs_ls, timeout = None ) ans = [] for item in results.choices: ans.append([item.text.strip()]) return ans def call_llama3_func( inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=20, temperature=0.0 ): print(max_decode_steps, temperature) output = call_llama3_single_prompt( inputs, model=model, max_decode_steps=max_decode_steps, temperature=temperature ) if isinstance(inputs, str): return output[0] else: return output def get_supervised_decom(queries_ls): prompts = [] for query in queries_ls: prompts.append([{"role": "user", "content": query}]) prompts_tokened = [tokenizer.apply_chat_template(x, tokenize=False, add_generation_prompt=True) for x in prompts] results = client.completions.create( model="supervised", max_tokens=512, temperature=0, prompt=prompts_tokened, timeout = None ) ans = [] for item in results.choices: ans.append([[x.strip() for x in item.text.split('|')]]) return ans def get_ans(queries_ls, passages): prompts = [] for i in range(len(queries_ls)): # Please answer the question to the best of your knowledge, even if the context does not directly provide the information. Use any relevant knowledge you have to provide a helpful answer. prompts.append([ {"role": "system", "content": "You are a helpful assistant. answer the question according to the context."}, {"role": "user", "content": 'Context: ' + '\n'.join(passages[i][:3])}, {"role": "user", "content": 'Question: ' + queries_ls[i]}, ]) pred_ls = [row[0] for row in call_llama3_func(prompts, max_decode_steps=100)] return pred_ls def hit_score(passages_ls, anspids, k): assert len(passages_ls) == len(anspids) cnt = 0 for i in range(len(passages_ls)): # 获取前k个检索结果 retrieved_topk = passages_ls[i][:k] # 将检索结果和答案都转换为小写以进行不区分大小写的匹配 retrieved_topk = [p.strip().lower() for p in retrieved_topk] ans_raw = anspids[i] # 检查每个可能的答案 for ans in ans_raw: ans = ans.strip().lower() # 如果答案在任何一个检索结果中出现 if any(ans in passage for passage in retrieved_topk): cnt += 1 break return cnt / len(passages_ls) def cover_em(pred_ls, ans_ls): assert len(pred_ls) == len(ans_ls) cnt = 0 for idx in range(len(pred_ls)): pred = pred_ls[idx].lower() ans = eval(ans_ls[idx]) for j in range(len(ans)): if ans[j].lower() in pred: cnt += 1 break return cnt/len(pred_ls) def supervised_method(): sub_query_str_l = get_supervised_decom(raw_queries) passages_ls = get_ir_result(raw_queries, sub_query_str_l) print(f"hit@20 :{hit_score(passages_ls, ans_pids, 20)}") print(f"hit@100 :{hit_score(passages_ls, ans_pids, 100)}") passages_ls = [item[:3] for item in passages_ls] pred_ls = get_ans(raw_queries, passages_ls) return cover_em(pred_ls, true_answers) def dense_method(): sub_query_str_l = [ [[raw]]for raw in raw_queries] passages_ls = get_ir_result(raw_queries, sub_query_str_l) print(f"hit@20 :{hit_score(passages_ls, ans_pids, 20)}") print(f"hit@100 :{hit_score(passages_ls, ans_pids, 100)}") passages_ls = [item[:3] for item in passages_ls] pred_ls = get_ans(raw_queries, passages_ls) return cover_em(pred_ls, true_answers) def unsupervised_method(): sub_query_str_l = [] for query in tqdm(raw_queries): url = 'http://127.0.0.1:50002/execute?query='+query response = requests.get(url=url) res_dic = response.json() sub_query_str_l.append([res_dic['text']]) passages_ls = get_ir_result(raw_queries, sub_query_str_l) print(f"hit@20 :{hit_score(passages_ls, ans_pids, 20)}") print(f"hit@100 :{hit_score(passages_ls, ans_pids, 100)}") passages_ls = [item[:3] for item in passages_ls] pred_ls = get_ans(raw_queries, passages_ls) return cover_em(pred_ls, true_answers) def iclfeed_method(): sub_query_str_l = [] for query in tqdm(raw_queries): url = 'http://127.0.0.1:50003/execute?query='+query response = requests.get(url=url) res_dic = response.json() sub_query_str_l.append([res_dic['text']]) passages_ls = get_ir_result(raw_queries, sub_query_str_l) print(f"hit@1 :{hit_score(passage_ls, ans_pids, 1)}") print(f"hit@2 :{hit_score(passage_ls, ans_pids, 2)}") print(f"hit@3 :{hit_score(passage_ls, ans_pids, 3)}") print(f"hit@20 :{hit_score(passages_ls, ans_pids, 20)}") print(f"hit@100 :{hit_score(passages_ls, ans_pids, 100)}") passages_ls = [item[:3] for item in passages_ls] pred_ls = get_ans(raw_queries, passages_ls) return cover_em(pred_ls, true_answers) def colbert_search(query_item): url = 'http://localhost:8895/api/search?query='+query_item+'&k=100' response = requests.get(url=url) res_dic = response.json() corpus_list_topk = res_dic['topk'] passage_ls = [] for i in range(100): passage_ls.append(corpus_list_topk[i]['text']) return passage_ls def colbert_method(): passage_ls = [] for query in tqdm(raw_queries): passage_ls.append(colbert_search(query)) print(f"hit@1 :{hit_score(passage_ls, ans_pids, 1)}") print(f"hit@2 :{hit_score(passage_ls, ans_pids, 2)}") print(f"hit@3 :{hit_score(passage_ls, ans_pids, 3)}") print(f"hit@20 :{hit_score(passage_ls, ans_pids, 20)}") print(f"hit@100 :{hit_score(passage_ls, ans_pids, 100)}") passage_ls = [item[:3] for item in passage_ls] pred_ls = get_ans(raw_queries, passage_ls) return cover_em(pred_ls, true_answers) if __name__ == '__main__': dataset_path = ragqa_paths.dataset_file(dataset, f"{dataset}_test.csv") tokenizer = AutoTokenizer.from_pretrained(ragqa_paths.LLAMA_MODEL) client = OpenAI(api_key="0",base_url="http://127.0.0.1:50001/v1") raw_data = pd.read_csv(dataset_path, header=0) raw_data = raw_data.drop_duplicates(subset=['question']) # raw_data = raw_data.head(10000) raw_queries = list(raw_data['question']) true_answers = list(raw_data['answers']) ans_pids = [eval(item) for item in list(raw_data['anspid'])] # print(supervised_method()) # print(unsupervised_method()) # print(iclfeed_method()) print(colbert_method()) # print(dense_method())