Spaces:
Sleeping
Sleeping
File size: 2,879 Bytes
e29422a |
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 |
from components.rag_components import get_length_without_special_tokens
import ipdb
QUERY_REWRITING = """Given a user query and its context (conversational history), decontextualize the question by addressing coreference and omission issues. The resulting question should retain its original meaning and be as informative as possible, and should not duplicate any previously asked questions in the context. JUST ANSWER THE RESOLVED QUERY WITHOUT ANY OTHER SENTENCES.\nContext: {}\n"""
REMINDER = """\nRemember you are a query rewriter. JUST ANSWER THE RESOLVED QUERY WITHOUT ANY OTHER SENTENCES."""
def get_context_from_message_history(message_history):
context = ""
for message in message_history:
if message["role"] not in ["system"]:
context += f'{message["role"]}: {message["content"]}\n'
return context if context else "No context available."
def rewrite_query(query: str, history: str, rewriter, rewriter_tokenizer, rewriter_terminator, device="cuda", max_tokens=256, temperature=0.0, top_p=0.9) -> str:
# ipdb.set_trace()
# DELETE LAST \n<assistant>\n
history = "\n".join(history.split("\n")[:-1])
system_prompt = QUERY_REWRITING.format(history)
query += REMINDER
user_prompt = [{"role": "system", "content": system_prompt}, {"role": "user", "content": f"user query: {query}"}]
prompt = rewriter_tokenizer.apply_chat_template(user_prompt, tokenize=False, add_generation_prompt=True)
print("user_prompt:", user_prompt)
print("PROMPT:", prompt)
# ipdb.set_trace()
print("System Prompt:", system_prompt)
print("Prompt:", prompt)
inputs = rewriter_tokenizer(prompt, return_tensors="pt").to(rewriter.device)
outputs = rewriter.generate(
**inputs,
max_new_tokens=max_tokens,
eos_token_id=rewriter_terminator,
do_sample=False, # Greedy decoding to be deterministic
# temperature=temperature,
top_p=top_p,
)
prompt_length = get_length_without_special_tokens(prompt, rewriter_tokenizer)
response = rewriter_tokenizer.decode(outputs[0], skip_special_tokens=True)[prompt_length:]
return response.strip()
# def rewrite_query(query: str, history: str, rewriter, rewriter_tokenizer, device="cuda") -> str:
# context = "|||".join([history, query])
# # rewriter = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device).eval()
# # rewriter_tokenizer = AutoTokenizer.from_pretrained(model_name)
# tokenized_context = rewriter_tokenizer.encode(context, return_tensors="pt").to(device)
# output_ids = rewriter.generate(
# tokenized_context,
# max_length=200,
# num_beams=4,
# repetition_penalty=2.5,
# length_penalty=1.0,
# early_stopping=True
# ).to(device)
# rewrite = rewriter_tokenizer.decode(output_ids[0], skip_special_tokens=True)
# return rewrite
|