# =================================================================
#  LoRA模型推理测试脚本 (首演)
# =================================================================

import torch
from unsloth import FastLanguageModel
from transformers import pipeline
import os

# --- 1. 配置路径 ---

# 您的Gemma基础模型的本地路径
base_model_path = "./gemma-3-4b-it-qat-unsloth-bnb-4bit"

# 您刚刚提取并重命名的LoRA模型的路径
# 请确保这个文件夹名和您自己保存的完全一致！
lora_model_path = "./lora_model_roleplayer_actor"

print("✅ 步骤 1/4: 路径配置完成。")
print(f"   - 基础模型路径: {base_model_path}")
print(f"   - LoRA模型路径: {lora_model_path}")

# --- 2. 加载模型与注入灵魂 ---

print("\n✅ 步骤 2/4: 正在加载基础模型...")

# 以4-bit精度加载基础模型，确保与训练时一致
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = base_model_path,
    load_in_4bit = True,
)

print("   - 基础模型加载成功。")

# 这是最关键的一步：将LoRA的“灵魂”注入到基础模型中！
# FastLanguageModel会自动处理所有复杂的合并工作
print(f"   - 正在从 '{lora_model_path}' 加载并注入LoRA适配器...")
model.load_adapter(lora_model_path)
print("   - LoRA灵魂注入成功！模型已准备就绪。")


# --- 3. 准备“剧本” (Prompt) ---

print("\n✅ 步骤 3/4: 正在准备测试剧本...")

# 我们使用和训练时完全相同的 Alpaca 格式
# 这能确保模型能最好地理解我们的指令
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{}

### Input:
{}

### Response:
{}"""

# --- 定义我们的角色和对话 ---

# 角色1: 薇莲 (一个我们训练过的角色)
instruction_virene = """角色名称：薇莲（Virene）
开场语：「真相，始终都存在于迷雾之中。」
身份背景：薇莲是一名神秘的赏金猎人，常常被人雇佣去完成各种危险任务，从而掩盖她本身的身份和目的。据传，薇莲早年曾在某个神秘组织中学习过各种神秘技能，所以她的能力非常高超。
性格特征：薇莲总是保持着冷静、沉着的态度，不论面对何种情况都能保持冷静。同时，她总是带有一定的神秘色彩，让人无法洞察她真正的想法和动机。她对任务非常认真，但很少会谈及自己的生活和过去，因此让人对她的身份感到好奇。
语言风格：薇莲的语言简洁有力，通常只说必要的话语来传达她的意思。她的语气总是带有一丝威慑力，让人不敢轻易挑战她。
行为特征：薇莲行动迅速而准确，总是在保持低调的同时完成任务。她具备很强的隐蔽能力，在执行任务的时候几乎不留痕迹，让人难以发现她的存在。不过，她也有时候会让人感到无法理解，经常出现在决定性瞬间，让人觉得她真正的动机仍旧是个谜。"""
input_virene = "我需要一个赏金猎人完成一个任务，听说您非常厉害。我们可以谈一下合作吗？"


# 角色2: （您可以自己创造一个新的角色来进行测试！）
instruction_new_role = """你现在是一个脾气火爆、说话直来直去，但内心充满正义感的退休老兵，名叫“老炮儿”。
你的口头禅是“嘿，我说你小子...”。
你的语言风格充满京味儿，简洁有力。"""
input_new_role = "大爷，问个路，这附近哪有吃饭的地方啊？"


# --- 选择一个角色进行测试 ---
# 您可以切换 instruction 和 input 来测试不同的角色
instruction_to_test = instruction_new_role
input_to_test = input_new_role

# 将“剧本”格式化
prompt = alpaca_prompt.format(
    instruction_to_test,
    input_to_test,
    "",  # Response部分留空，等待模型生成
)

print("   - 剧本已生成，准备开始推理！")

# --- 4. 开始推理 (见证奇迹的时刻) ---

print("\n✅ 步骤 4/4: 首演开始！正在生成对话...")

# 使用transformers的pipeline工具，这是最简单的推理方式
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer
)

# 设置生成参数
generation_args = {
    "max_new_tokens": 256,   # 最多生成多少个新词
    "do_sample": True,       # 开启采样，让回答更多样
    "temperature": 0.7,      # 温度，越低回答越稳定，越高越有创造力
    "top_p": 0.9,            # Top-p采样，控制多样性
    "top_k": 50,             # Top-k采样，控制多样性
    "pad_token_id": tokenizer.eos_token_id # 明确告知结束符
}

# 运行pipeline！
outputs = pipe(prompt, **generation_args)

# --- 打印结果 ---

print("\n\n==================== 🌟 演出结束 🌟 ====================")
print("完整的生成文本：\n")
print(outputs[0]['generated_text'])
print("\n========================================================")

# 只提取模型生成的部分，更清晰
response_part = outputs[0]['generated_text'].split("### Response:")[1].strip()
print("\n只看模型的回答部分：\n")
print(response_part)
print("\n========================================================")