reference/evaluation_drivers/test_webq.py

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())