Update handler.py

df67569 verified 6 months ago

2.03 kB

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}]

	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}]