reference/evaluation_drivers/test_manyqatext.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 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())