File size: 2,237 Bytes
e222084 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import subprocess
import sys
# Install VLLM using subprocess
subprocess.check_call([
f"{sys.executable}", "-m", "pip", "install",
"vllm @ https://github.com/vllm-project/vllm/releases/download/v0.6.1.post1/vllm-0.6.1.post1+cu118-cp310-cp310-manylinux1_x86_64.whl"
])
import json
import logging
import os
from vllm import LLM, SamplingParams
logger = logging.getLogger()
logger.setLevel(logging.INFO)
def model_fn(model_dir, context= None):
model = LLM(
model=model_dir,
trust_remote_code=True,
gpu_memory_utilization=0.9,
tensor_parallel_size=4 # Set tensor parallelism to 4 GPUs
)
return model
def predict_fn(data, model , context= None):
try:
input_text = data.pop("inputs", data)
parameters = data.pop("parameters", {})
# Qwen 2.5 chat template
chat_template = f"<|im_start|>system\nYou are a helpful AI assistant.<|im_end|>\n<|im_start|>user\n{input_text}<|im_end|>\n<|im_start|>assistant\n"
sampling_params = SamplingParams(
temperature=parameters.get("temperature", 0.7),
top_p=parameters.get("top_p", 0.9),
max_new_tokens=parameters.get("max_new_tokens", 512),
do_sample=True,
stop_tokens=["<|im_end|>", "<|im_start|>"] # Add stop tokens for Qwen
)
outputs = model.generate(chat_template, sampling_params)
generated_text = outputs[0].outputs[0].text
# Remove any trailing stop tokens if they were generated
for stop_token in ["<|im_end|>", "<|im_start|>"]:
if generated_text.endswith(stop_token):
generated_text = generated_text[:-len(stop_token)].strip()
return {"generated_text": generated_text}
except Exception as e:
logger.error(f"Exception during prediction: {e}")
return {"error": str(e)}
def input_fn(request_body, request_content_type,context= None):
if request_content_type == "application/json":
return json.loads(request_body)
else:
raise ValueError(f"Unsupported content type: {request_content_type}")
def output_fn(prediction, accept , context= None) :
return json.dumps(prediction) |