AntAngelMed

To run offline batched inference or an OpenAI-compatible API service on Nvidia A100 hardware, use vllm.

Environment Preparation

Install vllm version 0.11.0:

pip install vllm==0.11.0

Offline Inference Example

Use the following Python pattern to perform offline inference:

from modelscope import AutoTokenizer
from vllm import LLM, SamplingParams

def main():
    model_path = "MedAIBase/AntAngelMed" # model_id or your_local_model_path
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    sampling_params = SamplingParams(
        temperature=0.6,
        top_p=0.95,
        top_k=20,
        repetition_penalty=1.05,
        max_tokens=16384,
    )
    llm = LLM(
        model=model_path,
        trust_remote_code=True,
        dtype="bfloat16",
        tensor_parallel_size=4, 
    )

    prompt = "What should I do if I have a headache?" 
    messages = [
        {"role": "system", "content": "You are AntAngelMed, a helpfull medical assistant."},
        {"role": "user", "content": prompt},
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    outputs = llm.generate([text], sampling_params)
    print(outputs[0].outputs[0].text)
  
if __name__ == "__main__":
    main()

To run on Ascend NPUs, use the vllm-ascend backend. It is recommended to use an Ascend Atlas 800I A2 server (64*8GB memory version) and Docker for deployment.

Environment Preparation

1. Pull the Docker image:

docker pull quay.io/ascend/vllm-ascend:v0.11.0rc3

2. Run the container with necessary device mappings and volume mounts:

NAME=your_container_name
MODEL_PATH=your_absolute_model_path

docker run -itd --privileged --name=$NAME --net=host \
 --shm-size=1000g \
   --device /dev/davinci0 \
   --device /dev/davinci1 \
   --device /dev/davinci2\
   --device /dev/davinci3 \
   --device /dev/davinci4 \
   --device /dev/davinci5 \
   --device /dev/davinci6 \
   --device /dev/davinci7 \
   --device=/dev/davinci_manager \
   --device=/dev/hisi_hdc \
   --device /dev/devmm_svm \
   -v /usr/local/dcmi:/usr/local/dcmi \
    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
    -v /etc/ascend_install.info:/etc/ascend_install.info \
   -v /usr/local/sbin:/usr/local/sbin \
   -v /etc/hccn.conf:/etc/hccn.conf \
   -v $MODEL_PATH:$MODEL_PATH \
   quay.io/ascend/vllm-ascend:v0.11.0rc2 \
   bash

docker exec -u root -it $NAME bash

3. Configure required environment variables in the terminal:

export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export HCCL_OP_EXPANSION_MODE="AIV"
export NPU_MEMORY_FRACTION=0.97
export TASK_QUEUE_ENABLE=1
export OMP_NUM_THREADS=100
export ASCEND_LAUNCH_BLOCKING=0
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export VLLM_USE_MODELSCOPE=true

Offline Inference

from transformers import AutoTokenizer
from vllm import LLM, SamplingParams
model_path = "MedAIBase/AntAngelMed"
tokenizer = AutoTokenizer.from_pretrained(model_path)
sampling_params = SamplingParams(temperature=0.7, top_p=0.8, repetition_penalty=1.05, max_tokens=16384)
llm = LLM(model=model_path, 
    dtype='float16',
    tensor_parallel_size=4,                     
    gpu_memory_utilization=0.97,  
    enable_prefix_caching=True,  
    enable_expert_parallel=True,
    trust_remote_code=True)
prompt = "What should I do if I have a headache?"
messages = [
    {"role": "system", "content": "You are AntAngelMed, a helpful medical assistant."},
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
outputs = llm.generate([text], sampling_params)

Online Inference (OpenAI API Server)

model_id=MedAIBase/AntAngelMed
taskset -c 0-23 python3 -m vllm.entrypoints.openai.api_server \
  --model $model_id \
  --max-num-seqs=200 \
  --tensor-parallel-size 4 \
  --data-parallel-size 2 \
  --enable_expert_parallel \
  --gpu_memory_utilization 0.97 \
  --served-model-name AntAngelMed \
  --max-model-len 32768 \
  --port 8080 \
  --enable-prefix-caching \
  --block-size 128 \
  --async-scheduling \
  --trust_remote_code

Query the server:

curl http://0.0.0.0:8080/v1/chat/completions -d '{
    "model": "AntAngelMed",
    "messages": [
        {
            "role": "system",
            "content": "You are a helpful assistant."
        },
        {
            "role": "user",
            "content": "What should I do if I have a headache?"
        }
    ],
    "temperature": 0.6
}'

To use the chat model locally, use the transformers library. Ensure you set trust_remote_code=True because the model utilizes a custom architecture (Ling-flash-2.0 MoE). The model supports a 128K context length via YaRN extrapolation.

from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "MedAIBase/AntAngelMed"  # model_id or your_local_model_path
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "What should I do if I have a headache?"
messages = [
    {"role": "system", "content": "You are AntAngelMed, a helpful medical assistant."},
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt", return_token_type_ids=False).to(model.device)
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=16384,
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)