import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

class EndpointHandler:
    def __init__(self, path=""):
        # 🚨 配置信息：指向你的基座和其中一个挂件
        self.base_model_id = "Qwen/Qwen3.5-27B"
        self.lora_id = "LucasUn/ASQ-Hypotheses"
        
        print(f"🚀 [Native Mode] 正在加载基座模型 (BF16): {self.base_model_id}")
        self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_id, trust_remote_code=True)
        
        # 1. 加载基座 (在 A100 80G 上，这步约占用 54GB 显存)
        base_model = AutoModelForCausalLM.from_pretrained(
            self.base_model_id,
            torch_dtype=torch.bfloat16,
            device_map="auto", # 自动分发到 A100
            trust_remote_code=True
        )
        
        print(f"🔗 [Native Mode] 正在物理合并 LoRA 插件: {self.lora_id}")
        # 2. 加载 LoRA 插件
        self.model = PeftModel.from_pretrained(base_model, self.lora_id)
        
        # 3. 核心优化：物理合并权重。合并后推理速度等同于无 LoRA 状态，且极致稳定
        self.model = self.model.merge_and_unload()
        self.model.eval()
        
        print("✅ [Native Mode] 模型已就绪，所有索引匹配完美！")

    def __call__(self, data):
        inputs = data.pop("inputs", data)
        
        # 编码输入
        input_ids = self.tokenizer(inputs, return_tensors="pt").to("cuda")
        
        # 执行生成
        with torch.no_grad():
            output_ids = self.model.generate(
                **input_ids,
                max_new_tokens=512,
                temperature=0.3,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id
            )
        
        # 解码
        response = self.tokenizer.decode(
            output_ids[0][input_ids.input_ids.shape[1]:], 
            skip_special_tokens=True
        )
        return [{"generated_text": response}]