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 manyqa_text 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/' # os.environ['CUDA_VISIBLE_DEVICES'] = '1' dataset = 'manyqa_text' def call_llama3_single_prompt( inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=100, 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=100, 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. # "You are a helpful assistant. answer the question according to the context." "You are a helpful assistant. 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)): # print(len(passages_ls[i])) # 获取前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(passage in ans for passage in retrieved_topk): cnt += 1 break return cnt / len(passages_ls) def hit_ls(passages_ls, anspids, k): assert len(passages_ls) == len(anspids) ans_ls = [] for i in range(len(passages_ls)): hit = 0 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): hit = 1 break ans_ls.append(hit) return ans_ls import re import csv from io import StringIO def parse_string(s): s = s.strip() if s.startswith('[') and s.endswith(']'): # 处理列表结构 inner = s[1:-1].strip() # 将单引号包裹的元素替换为双引号包裹，并转义内部双引号 pattern = re.compile(r"'((?:[^'\\]|\\.)*?)'") def replace(match): content = match.group(1) content = content.replace('"', r'\"') return f'"{content}"' new_inner = pattern.sub(replace, inner) # 使用 csv.reader 解析处理后的内容 csv_reader = csv.reader( StringIO(new_inner), quotechar='"', escapechar='\\', skipinitialspace=True ) try: return next(csv_reader) except StopIteration: return [] else: # 处理单个字符串，包裹为双引号并转义内部双引号 content = s.replace('"', r'\"') csv_reader = csv.reader( StringIO(f'"{content}"'), quotechar='"', escapechar='\\' ) try: return next(csv_reader) except StopIteration: return [s] def cover_em(pred_ls, ans_ls): assert len(pred_ls) == len(ans_ls) cnt = 0 score_ls = [] # 用于记录每个查询的正确性 for idx in range(len(pred_ls)): pred = pred_ls[idx].lower() ans = str(ans_ls[idx]) correct = 0 # 默认为错误 if ans.lower() in pred: cnt += 1 correct = 1 score_ls.append(correct) return cnt/len(pred_ls), score_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[: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@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[: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 = sentence_search(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[:3] 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, "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) passages_ls = sentence_search(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[:3] for item in passages_ls] pred_ls = get_ans(raw_queries, passages_ls) return cover_em(pred_ls, true_answers) def sentence_search(queries, questions_ls): passage_ls = [] for idx in tqdm(range(len(questions_ls))): top_100_passages = [] query = questions_ls[idx] payload = { "query" : queries[idx], "sub_query": query, "k": 100 } response = requests.post( 'http://localhost:8220/api/search', json=payload ) for i in range(len(response.json()['topk'])): top_100_passages.append(response.json()['topk'][i]['text']) passage_ls.append(top_100_passages) return passage_ls 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'] # print(len(corpus_list_topk)) passage_ls = [] for i in range(min(100, len(corpus_list_topk))): 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) score, score_ls = cover_em(pred_ls, true_answers) from datetime import datetime timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') log_data = { 'raw_query': raw_queries, 'true_answer': true_answers, 'predicted_answer': pred_ls, 'acc' : score_ls, 'ir_hit' : hit_ls(passage_ls, ans_pids, 1) } log_df = pd.DataFrame(log_data) log_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs", f"colbert_method_log_{timestamp}.csv") log_df.to_csv(log_path, index=False, encoding='utf-8-sig') print(f"日志已保存至: {log_path}") return score 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@2 :{hit_score(passages_ls, ans_pids, 2)}") 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)[0] 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']) # raw_queries = [q.strip() + '?' if not q.strip().endswith('?') else q.strip() for q in raw_data['question']] # raw_queries = [q if not q.endswith('?') else q.strip('?') for q in raw_data['question']] true_answers = list(raw_data['answers']) ans_pids = [[item] for item in list(raw_data['anspid'])] # print(supervised_method()) # print(unsupervised_method()) # print(iclfeed_method()) # print(colbert_method()) # print(dense_method()) print(sentence_colbert())