reference/evaluation_drivers/test_triviaqa.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 = '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())