Create inference.py

inference.py

@@ -0,0 +1,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)