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_nq import * import requests def get_vllm_llama(temperature, max_tokens, chats): url = "http://127.0.0.1:60000/ask" data = { "temperature": temperature, "max_tokens": max_tokens, "chats": chats, } response = requests.post(url, json=data, timeout=None) response_data = response.json() passage_ls = response_data.get("output", []) return passage_ls def call_llama3_single_prompt( inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=20, temperature=0.8 ): 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) return ans def call_llama3_func( inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=1024, 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": "system", "content": "You are a query decomposition assistant. Please decompose one query Q into semantically coherent sub-queries."}, {"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 ) ans = [] for item in results.choices: ans.append([[x.strip() for x in item.text.split('|')]]) return ans def get_lora_ans(queries_ls, passages): prompts = [] for i in range(len(queries_ls)): prompts.append([ {"role": "system", "content": "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."}, {"role": "user", "content": 'Context: ' + '\n'.join(passages[i])}, {"role": "user", "content": 'Question: ' + queries_ls[i]}, ]) prompts_tokened = [tokenizer.apply_chat_template(x, tokenize=False, add_generation_prompt=True) for x in prompts] results = client.completions.create( model="web_ft", max_tokens=100, temperature=0, prompt=prompts_tokened ) ans = [] for item in results.choices: ans.append(item.text.strip()) return ans def get_ans(queries_ls, passages): prompts = [] for i in range(len(queries_ls)): prompts.append([ {"role": "system", "content": "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."}, {"role": "user", "content": 'Context: ' + '\n'.join(passages[i])}, {"role": "user", "content": 'Question: ' + queries_ls[i]}, ]) prompts_tokened = [tokenizer.apply_chat_template(x, tokenize=False, add_generation_prompt=True) for x in prompts] # pred_ls = [row[0] for row in call_llama3_func(prompts, max_decode_steps=100)] results = client.completions.create( model=ragqa_paths.LLAMA_MODEL, max_tokens=100, temperature=0, prompt=prompts_tokened ) ans = [] for item in results.choices: ans.append(item.text.strip()) return ans 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) 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']]) print(sub_query_str_l) passages_ls = get_ir_result(raw_queries, sub_query_str_l) 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) pred_ls = get_ans(raw_queries, passages_ls) return cover_em(pred_ls, true_answers) if __name__ == '__main__': dataset_path = ragqa_paths.dataset_file("nq", "nq_test.csv") tokenizer = AutoTokenizer.from_pretrained(ragqa_paths.LLAMA_MODEL) client = OpenAI(api_key="0",base_url="http://0.0.0.0:50001/v1") raw_data = pd.read_csv(dataset_path, header=0) # raw_data = raw_data.head(10) raw_queries = list(raw_data['question']) true_answers = list(raw_data['answers']) print(supervised_method()) print(unsupervised_method()) print(iclfeed_method()) # with open("/root/autodl-tmp/result.txt", "a") as file: # file.write(f"webq_supervised_method_{supervised_method()}\n") # file.write(f"webq_unsupervised_method_{unsupervised_method()}\n") # file.write(f"webq_iclfeed_method_{iclfeed_method()}\n")