111 lines
5.0 KiB
Python
111 lines
5.0 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from typing import List, Tuple, Dict, Any
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
from OHS.Script import *
|
|
|
|
class Reranker:
|
|
def __init__(self, model_dir: str, device):
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_dir, padding_side='left', use_fast=True)
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
model_dir,
|
|
dtype="bfloat16",
|
|
attn_implementation="flash_attention_2"
|
|
).eval()
|
|
|
|
self.model.to(device)
|
|
|
|
self.token_false_id = self.tokenizer.convert_tokens_to_ids("no")
|
|
self.token_true_id = self.tokenizer.convert_tokens_to_ids("yes")
|
|
self.max_length = 8192
|
|
|
|
self.prefix = "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n<|im_start|>user\n"
|
|
self.suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
|
|
self.prefix_tokens = self.tokenizer.encode(self.prefix, add_special_tokens=False)
|
|
self.suffix_tokens = self.tokenizer.encode(self.suffix, add_special_tokens=False)
|
|
|
|
self.task_description4document = "给定客户异议，根据话术的目标解决问题和目标客户特征，判断该异议处理话术是否能够解决该异议。"
|
|
self.task_description4script = "给定客户异议和相关处理话术，判断该异议处理话术是否能够解决该异议。"
|
|
|
|
def format_instruction4document(self, query: str, document: str) -> str:
|
|
return f"<Instruct>: {self.task_description4document}\n<Query>: {query}\n<Document>: {document}"
|
|
|
|
def format_instruction4script(self, query: str, document: str) -> str:
|
|
return f"<Instruct>: {self.task_description4script}\n<Query>: {query}\n<Document>: {document}"
|
|
|
|
def _process_inputs(self, texts: List[str]) -> Dict[str, torch.Tensor]:
|
|
inputs = self.tokenizer(
|
|
texts,
|
|
padding=False,
|
|
truncation='longest_first',
|
|
return_attention_mask=False,
|
|
max_length=self.max_length - len(self.prefix_tokens) - len(self.suffix_tokens)
|
|
)
|
|
|
|
for i, ele in enumerate(inputs['input_ids']):
|
|
inputs['input_ids'][i] = self.prefix_tokens + ele + self.suffix_tokens
|
|
|
|
inputs = self.tokenizer.pad(inputs, padding=True, return_tensors="pt", max_length=self.max_length)
|
|
|
|
for key in inputs:
|
|
inputs[key] = inputs[key].to(self.model.device)
|
|
|
|
return inputs
|
|
|
|
@torch.no_grad()
|
|
def compute_scores4document(self, query: str, documents: List[str]) -> List[float]:
|
|
pairs = [self.format_instruction4document(query, doc) for doc in documents]
|
|
inputs = self._process_inputs(pairs)
|
|
|
|
batch_scores = self.model(**inputs).logits[:, -1, :]
|
|
|
|
true_vector = batch_scores[:, self.token_true_id]
|
|
false_vector = batch_scores[:, self.token_false_id]
|
|
batch_scores = torch.stack([false_vector, true_vector], dim=1)
|
|
batch_scores = F.log_softmax(batch_scores, dim=1)
|
|
scores = batch_scores[:, 1].exp().tolist()
|
|
|
|
return scores
|
|
|
|
@torch.no_grad()
|
|
def compute_scores4script(self, query: str, documents: List[str]) -> List[float]:
|
|
pairs = [self.format_instruction4script(query, doc) for doc in documents]
|
|
inputs = self._process_inputs(pairs)
|
|
|
|
batch_scores = self.model(**inputs).logits[:, -1, :]
|
|
|
|
true_vector = batch_scores[:, self.token_true_id]
|
|
false_vector = batch_scores[:, self.token_false_id]
|
|
batch_scores = torch.stack([false_vector, true_vector], dim=1)
|
|
batch_scores = F.log_softmax(batch_scores, dim=1)
|
|
scores = batch_scores[:, 1].exp().tolist()
|
|
|
|
return scores
|
|
|
|
def rerank(self, query: str, top_scripts: List[ObjectionHandleScript], rerank_N: int = 10) -> List[Tuple[ObjectionHandleScript, float]]:
|
|
"""
|
|
重排序函数
|
|
Args:
|
|
query: 查询文本
|
|
top_scripts: 召回的话术列表，每个元素为话术对象
|
|
Returns:
|
|
重排序后的话术列表，每个元素为(话术对象, 重排序分数)
|
|
"""
|
|
if not top_scripts:
|
|
return []
|
|
scripts = [script.get_script() for script in top_scripts]
|
|
documents = [script.get_script() for script in top_scripts]
|
|
|
|
document_scores = self.compute_scores4document(query, documents)
|
|
script_scores = self.compute_scores4script(query, scripts)
|
|
|
|
results = []
|
|
for i, (script, doc_score, scr_score) in enumerate(zip(top_scripts, document_scores, script_scores)):
|
|
final_score = 0.6 * doc_score + 0.4 * scr_score
|
|
results.append((script, final_score))
|
|
|
|
results.sort(key=lambda x: x[1], reverse=True)
|
|
results = results[:rerank_N]
|
|
return results
|