from vllm import LLM, SamplingParams
import math

# 使用vLLM初始化模型
model = LLM(
    model="jhu-clsp/rank1-3b",
    tensor_parallel_size=1,  # GPU数量
    trust_remote_code=True,
    max_model_len=16000,     # 上下文长度
    gpu_memory_utilization=0.9,
    dtype="float16",
)

# 设置采样参数
sampling_params = SamplingParams(
    temperature=0,
    max_tokens=8192,
    logprobs=20,
    stop=["</think> true", "</think> false"],
    skip_special_tokens=False
)

# 准备提示词
def create_prompt(query, document):
    return (
        "确定以下段落是否与查询相关。"
        "仅回答'true'或'false'。\n"
        f"查询: {query}\n"
        f"段落: {document}\n"
        "<think>"
    )

# 示例用法
query = "气候变化有什么影响?"
document = "气候变化导致海平面上升、极端天气事件和生态系统破坏。这些影响是由人类活动导致大气中温室气体浓度增加引起的。"

# 生成预测
prompt = create_prompt(query, document)
outputs = model.generate([prompt], sampling_params)

# 提取分数
output = outputs[0].outputs[0]
text = output.text
final_logits = output.logprobs[-1]

# 获取"true"和"false"标记的ID
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/rank1-3b")
true_token = tokenizer(" true", add_special_tokens=False).input_ids[0]
false_token = tokenizer(" false", add_special_tokens=False).input_ids[0]

# 计算相关性分数("true"的概率)
true_logit = final_logits[true_token].logprob
false_logit = final_logits[false_token].logprob
true_score = math.exp(true_logit)
false_score = math.exp(false_logit)
relevance_score = true_score / (true_score + false_score)

print(f"推理链: {text}")
print(f"相关性分数: {relevance_score}")