File size: 5,277 Bytes
cd6775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | from flask import Flask, render_template, request, jsonify
from functools import lru_cache
import math
import os
from dotenv import load_dotenv
from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Searcher
load_dotenv()
INDEX_NAME = os.getenv("INDEX_NAME")
INDEX_ROOT = os.getenv("INDEX_ROOT")
app = Flask(__name__)
# #searcher = Searcher(index=f"{INDEX_ROOT}/{INDEX_NAME}")
# searcher = Searcher(index=f"/home/icml01/multi_rag/RAG/Search-in-the-Chain/ColBERT/experiments/strategyqa/indexes/strategyqa.nbits=2")
import sys
sys.path.append(r"/home/icml01/multi_rag/RAG/Decompose_retrieval")
from multihop_ir import *
from transformers import AutoTokenizer
from openai import OpenAI
instruction = 'You are a query decomposition assistant. Please decompose one query Q into semantically coherent sub-queries.'
model_id = "/home/icml01/Models/Llama-3.1-8B-Instruct"
pipe = pipeline(
"text-generation",
model=model_id,
model_kwargs={"torch_dtype": torch.bfloat16},
device="cuda:1",
)
def call_llama3_single_prompt(
input_str, model="Llama-3.1-8B-Instruct", max_decode_steps=200, temperature=0.8
):
if isinstance(input_str, str):
messages = [
{"role": "user", "content": input_str},
]
else:
messages = input_str
if temperature > 0:
outputs = pipe(
messages,
max_new_tokens=max_decode_steps,
temperature=temperature,
pad_token_id=pipe.tokenizer.eos_token_id,
)
else:
outputs = pipe(
messages,
max_new_tokens=max_decode_steps,
do_sample = False,
pad_token_id=pipe.tokenizer.eos_token_id,
)
return outputs
def call_llama3_func(
inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=200, temperature=0.0
):
outputs = []
# for input_str in inputs:
output = call_llama3_single_prompt(
inputs,
model=model,
max_decode_steps=max_decode_steps,
temperature=temperature
)
for item in output:
outputs.append([item[0]["generated_text"][-1]["content"]])
return outputs
def gen_prompt(# 分解结果产生
query,
instruction):
prompt = []
if instruction:
prompt.append({"role": "system", "content": instruction})
prompt.append({"role": "system", "content":"Instruction: Split the question into key fragments separated by |. Do not generate sub-queries, rewrite the question, or include explanations. For example, if the input is 'What color is the Santa Anita Park logo?', output 'Santa Anita Park| logo'. Generating sub-queries like 'What color is the logo?' is incorrect."})
# prompt.append({'role': 'user','content':'What color is the Santa Anita Park logo?'})
# prompt.append({'role': 'assistant','content':'Santa Anita Park| logo'})
prompt.append({"role": "user", "content": query})
return prompt
def get_sub_query(query):
assert isinstance(query,str)
message = [gen_prompt(query, instruction)]
print(message)
tmp_ls = call_llama3_func(message)[0][0].replace("\n", "").split('|')
# sub_q_ls = [[list(set([item.strip() for item in tmp_ls if item.strip() and item.strip() in query]))]]
sub_q_ls = [[list(set([item.strip() for item in tmp_ls]))]]
return sub_q_ls
counter = {"api" : 0}
# @lru_cache(maxsize=1000000)
# def api_search_query(query, k):
# print(f"Query={query}")
# if k == None: k = 10
# k = min(int(k), 100)
# pids, ranks, scores = searcher.search(query, k=100)
# pids, ranks, scores = pids[:k], ranks[:k], scores[:k]
# passages = [searcher.collection[pid] for pid in pids]
# probs = [math.exp(score) for score in scores]
# probs = [prob / sum(probs) for prob in probs]
# topk = []
# for pid, rank, score, prob in zip(pids, ranks, scores, probs):
# text = searcher.collection[pid]
# d = {'text': text, 'pid': pid, 'rank': rank, 'score': score, 'prob': prob}
# topk.append(d)
# topk = list(sorted(topk, key=lambda p: (-1 * p['score'], p['pid'])))
# return {"query" : query, "topk": topk}
# @lru_cache(maxsize=1000000)
def api_search_query(query, k):
print(f"Query={query}")
sub_query_str_l = get_sub_query(query)
print(sub_query_str_l)
return {"text" : get_ir_result(query, sub_query_str_l)}
@app.route("/api/update_instruction", methods=["POST"])
def update_instruction():
global instruction # 使用 global 来更新全局的 instruction 变量
# 从请求的 JSON 数据中获取新的 instruction
data = request.get_json()
if "instruction" in data:
instruction = data["instruction"]
return jsonify({"message": "Instruction updated successfully.", "new_instruction": instruction}), 200
else:
return jsonify({"message": "Instruction not provided in the request."}), 400
@app.route("/api/search", methods=["GET"])
def api_search():
if request.method == "GET":
counter["api"] += 1
print("API request count:", counter["api"])
return api_search_query(request.args.get("query"), request.args.get("k"))
else:
return ('', 405)
if __name__ == "__main__":
app.run("0.0.0.0", 8894)
|