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 = 'webq' model_name = ragqa_paths.LLAMA_LORA_WEBQ import re ITEM = 1 def filter_subqueries(data, queries): filtered_data = [] for group, query in zip(data, queries): filtered_group = [] query_tokens = set(re.findall(r"\w+", query.lower())) # 提取queries中的单词 for subquery in group[0]: subquery_tokens = re.findall(r"\w+", subquery) # 提取subquery中的单词 filtered_subquery = " ".join([token for token in subquery_tokens if token.lower() in query_tokens]) filtered_group.append(filtered_subquery) filtered_data.append([filtered_group]) return filtered_data 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=model_name, # 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 gen_prompt(): prompts = [] for query in raw_queries: prompt = [] prompt.append({"role": "system", "content": """"Given the input query, break it down into meaningful tokens like ColBERT. Ensure the tokens retain the key semantic components. Provide the output as a comma-separated list. Query: '{query}' Tokens:"""}) prompt.append({"role": "user", "content": "Query: Victoria Hong Kong has many what type of buildings?"}) prompt.append({"role": "assistant", "content": "Victoria, Hong, Kong, has, many, what, type, of, buildings,?"}) prompt.append({"role": "user", "content": f"Query: {query}" }) prompts.append(prompt) return prompts def sentence_colbert(): tmp = gen_prompt() raw_op_prompts = call_llama3_func(tmp) sub_queries_ls= [] for idx in range(len(raw_queries)): tmp_ls = raw_op_prompts[idx][0].replace("\n", "").split(",") tmp_ls = [list(set([item.strip() for item in tmp_ls if item.strip() and item.strip() in raw_queries[idx]]))] if len(tmp_ls[0]) == 0: tmp_ls = [[raw_queries[idx]]] sub_queries_ls.append(tmp_ls) passages_ls = get_ir_result(raw_queries, sub_queries_ls) print(f"hit@1 :{hit_score(passages_ls, ans_pids, 1)}") print(f"hit@3 :{hit_score(passages_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[:ITEM] for item in passages_ls] pred_ls = get_ans(raw_queries, passages_ls) return cover_em(pred_ls, true_answers) 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=ragqa_paths.LLAMA_LORA_WEBQ, 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][:ITEM])}, {"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 get_woans(queries_ls): 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."}, {"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 hit_score(passages_ls, anspids, k): # assert len(passages_ls) == len(anspids) # cnt = 0 # for i in range(len(passages_ls)): # retrieved_100_topk = passages_ls[i][:k] # ans_raw = anspids[i] # for ans in ans_raw: # if ans in retrieved_100_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(): cache_path = ragqa_paths.dataset_file(dataset, "s_cache.pkl") if os.path.exists(cache_path): with open(cache_path, 'rb') as f: sub_query_str_l = pickle.load(f) print("Loaded cached subqueries") else: sub_query_str_l = get_supervised_decom(raw_queries) with open(cache_path, 'wb') as f: pickle.dump(sub_query_str_l, f) print("Saved new subqueries to cache") passages_ls = get_ir_result(raw_queries, sub_query_str_l) print(f"hit@1 :{hit_score(passages_ls, ans_pids, 1)}") print(f"hit@2 :{hit_score(passages_ls, ans_pids, 2)}") print(f"hit@3 :{hit_score(passages_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[:ITEM] 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@1 :{hit_score(passages_ls, ans_pids, 1)}") print(f"hit@2 :{hit_score(passages_ls, ans_pids, 2)}") print(f"hit@3 :{hit_score(passages_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[:ITEM] for item in passages_ls] pred_ls = get_ans(raw_queries, passages_ls) return cover_em(pred_ls, true_answers) def wo_method(): pred_ls = get_woans(raw_queries) return cover_em(pred_ls, true_answers) def unsupervised_method(): cache_path = ragqa_paths.dataset_file(dataset, "un_cache.pkl") if os.path.exists(cache_path): with open(cache_path, 'rb') as f: sub_query_str_l = pickle.load(f) print("Loaded cached subqueries") else: 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']]) with open(cache_path, 'wb') as f: pickle.dump(sub_query_str_l, f) print("Saved new subqueries to cache") passages_ls = get_ir_result(raw_queries, sub_query_str_l) print(f"hit@1 :{hit_score(passages_ls, ans_pids, 1)}") print(f"hit@2 :{hit_score(passages_ls, ans_pids, 2)}") print(f"hit@3 :{hit_score(passages_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[:ITEM] for item in passages_ls] pred_ls = get_ans(raw_queries, passages_ls) return cover_em(pred_ls, true_answers) def iclfeed_method(): cache_path = ragqa_paths.dataset_file(dataset, "gpt_icl_cache.pkl") if os.path.exists(cache_path): with open(cache_path, 'rb') as f: sub_query_str_l = pickle.load(f) print("Loaded cached subqueries") else: 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']]) # 保存生成的子查询 with open(cache_path, 'wb') as f: pickle.dump(sub_query_str_l, f) print("Saved new subqueries to cache") passages_ls = get_ir_result(raw_queries, sub_query_str_l) print(f"hit@1 :{hit_score(passages_ls, ans_pids, 1)}") print(f"hit@2 :{hit_score(passages_ls, ans_pids, 2)}") print(f"hit@3 :{hit_score(passages_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[:ITEM] 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:8896/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[:ITEM] 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(model_name) client = OpenAI(api_key="0",base_url="http://127.0.0.1:50001/v1") client1 = OpenAI(api_key="0",base_url="http://127.0.0.1:20001/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(sentence_colbert()) # print(dense_method()) # print(wo_method())