|
import json |
|
import subprocess |
|
import torch |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
from typing import List, Dict |
|
|
|
|
|
subprocess.check_call(['pip', 'install', 'accelerate']) |
|
|
|
from accelerate import init_empty_weights, infer_auto_device_map, dispatch_model |
|
|
|
|
|
def format_chat(messages: List[Dict[str, str]], tokenizer) -> str: |
|
""" |
|
Format chat messages using Qwen's chat template |
|
""" |
|
return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
|
|
|
|
def model_fn(model_dir,context=None): |
|
""" |
|
Load the model and tokenizer from the model directory for inference. |
|
This version supports tensor parallelism across 4 GPUs. |
|
""" |
|
device_map = "auto" |
|
|
|
|
|
with init_empty_weights(): |
|
model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype="auto") |
|
|
|
|
|
device_map = infer_auto_device_map(model, max_memory={i: "25GiB" for i in range(4)}, no_split_module_classes=["QwenBlock"]) |
|
|
|
|
|
model = dispatch_model(model, device_map=device_map, offload_buffers=True) |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_dir) |
|
|
|
return model, tokenizer |
|
|
|
|
|
def predict_fn(input_data, model_and_tokenizer): |
|
""" |
|
Generate predictions for the input data. |
|
""" |
|
try: |
|
model, tokenizer = model_and_tokenizer |
|
data = json.loads(input_data) |
|
|
|
|
|
messages = data.get("messages", []) |
|
formatted_prompt = format_chat(messages, tokenizer) |
|
|
|
|
|
inputs = tokenizer([formatted_prompt], return_tensors="pt").to("cuda:0") |
|
|
|
|
|
outputs = model.generate( |
|
inputs['input_ids'], |
|
max_new_tokens=data.get("max_new_tokens", 512), |
|
temperature=data.get("temperature", 0.7), |
|
top_p=data.get("top_p", 0.9), |
|
repetition_penalty=data.get("repetition_penalty", 1.0), |
|
length_penalty=data.get("length_penalty", 1.0), |
|
do_sample=True |
|
) |
|
|
|
|
|
generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0] |
|
|
|
|
|
response = { |
|
"id": "chatcmpl-uuid", |
|
"object": "chat.completion", |
|
"model": "qwen-72b", |
|
"choices": [{ |
|
"index": 0, |
|
"message": { |
|
"role": "assistant", |
|
"content": generated_text |
|
}, |
|
"finish_reason": "stop" |
|
}], |
|
"usage": { |
|
"prompt_tokens": len(inputs['input_ids'][0]), |
|
"completion_tokens": len(outputs[0]), |
|
"total_tokens": len(inputs['input_ids'][0]) + len(outputs[0]) |
|
} |
|
} |
|
return response |
|
|
|
except Exception as e: |
|
return {"error": str(e), "details": repr(e)} |
|
|
|
|
|
def input_fn(serialized_input_data, content_type): |
|
""" |
|
Prepare the input data for inference. |
|
""" |
|
return serialized_input_data |
|
|
|
|
|
def output_fn(prediction_output, accept): |
|
""" |
|
Convert the model output to a JSON response. |
|
""" |
|
return json.dumps(prediction_output) |
|
|