56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
|
prompt = "你好，我的名字是"
|
|||
|
|
model_path = './model_ckpt/hairuo'
|
|||
|
|
# run_type = 'vllm'
|
|||
|
|
run_type = 'transformers'
|
|||
|
|
|
|||
|
|
if run_type == 'transformers':
|
|||
|
|
from ihp.zoo.hairuo import HairuoTokenizer
|
|||
|
|
from ihp.zoo.hairuo import HairuoForCausalLM
|
|||
|
|
|
|||
|
|
model = HairuoForCausalLM.from_pretrained(model_path)
|
|||
|
|
tokenizer = HairuoTokenizer.from_pretrained(model_path)
|
|||
|
|
|
|||
|
|
model.requires_grad_(False)
|
|||
|
|
model.eval()
|
|||
|
|
|
|||
|
|
inputs = tokenizer(prompt, return_tensors="pt")
|
|||
|
|
generate_ids = model.generate(inputs.input_ids, attention_mask = inputs.attention_mask, max_length=200, temperature=0.8, do_sample=True, eos_token_id=151644, pad_token_id=151644)
|
|||
|
|
generated_text = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
|
|||
|
|
print(generated_text)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if run_type == 'vllm':
|
|||
|
|
|
|||
|
|
# 载入 LLM 和 SamplingParams
|
|||
|
|
from vllm import LLM, SamplingParams
|
|||
|
|
from vllm import ModelRegistry
|
|||
|
|
from ihp.zoo.hairuo.vllm_hairuo import HairuoForCausalLM
|
|||
|
|
ModelRegistry.register_model("HairuoForCausalLM", HairuoForCausalLM)
|
|||
|
|
# 推理数据以List[str]格式组织
|
|||
|
|
prompts = [
|
|||
|
|
"你好，我的名字是",
|
|||
|
|
"The president of the United States is",
|
|||
|
|
"The capital of France is",
|
|||
|
|
"AI的未来是什么？",
|
|||
|
|
]
|
|||
|
|
# 设置采样参数
|
|||
|
|
sampling_params = SamplingParams(temperature=0.8, top_p=1)
|
|||
|
|
# 加载模型
|
|||
|
|
llm = LLM(
|
|||
|
|
model=model_path,
|
|||
|
|
trust_remote_code=True,
|
|||
|
|
tensor_parallel_size=1,
|
|||
|
|
# dtype='float32',
|
|||
|
|
gpu_memory_utilization=0.95,
|
|||
|
|
max_model_len=100,
|
|||
|
|
enforce_eager=True,
|
|||
|
|
)
|
|||
|
|
# 执行推理
|
|||
|
|
outputs = llm.generate(prompts, sampling_params)
|
|||
|
|
|
|||
|
|
# 输出推理结果
|
|||
|
|
for output in outputs:
|
|||
|
|
prompt = output.prompt
|
|||
|
|
generated_text = output.outputs[0].text
|
|||
|
|
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|