[update] fp16 inference

pull/825/head
jingyaogong 2026-03-27 16:29:46 +08:00
parent deed5674c4
commit 8f3169dba2
4 changed files with 4 additions and 4 deletions

View File

@ -27,7 +27,7 @@ def init_model(args):
else: else:
model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
get_model_params(model, model.config) get_model_params(model, model.config)
return model.eval().to(args.device), tokenizer return model.half().eval().to(args.device), tokenizer
def main(): def main():
parser = argparse.ArgumentParser(description="MiniMind模型推理与对话") parser = argparse.ArgumentParser(description="MiniMind模型推理与对话")

View File

@ -64,7 +64,7 @@ def init_model(args):
else: else:
model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
get_model_params(model, model.config) get_model_params(model, model.config)
return model.eval().to(args.device), tokenizer return model.half().eval().to(args.device), tokenizer
def parse_tool_calls(text): def parse_tool_calls(text):

View File

@ -44,7 +44,7 @@ def init_model(args):
else: else:
model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
print(f'MiniMind模型参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M(illion)') print(f'MiniMind模型参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M(illion)')
return model.eval().to(device), tokenizer return model.half().eval().to(device), tokenizer
class ChatRequest(BaseModel): class ChatRequest(BaseModel):

View File

@ -205,7 +205,7 @@ def load_model_tokenizer(model_path):
model_path, model_path,
trust_remote_code=True trust_remote_code=True
) )
model = model.eval().to(device) model = model.half().eval().to(device)
return model, tokenizer return model, tokenizer