Update inference.py

#6
by FINGU-AI - opened
Files changed (1) hide show
  1. inference.py +11 -1
inference.py CHANGED
@@ -1,9 +1,19 @@
1
  import os
2
  import json
 
 
3
  import torch
 
 
 
 
 
 
 
 
 
4
  from vllm import LLM, SamplingParams
5
  from vllm.utils import random_uuid
6
- from typing import List, Dict
7
 
8
  # Function to format chat messages using Qwen's chat template
9
  def format_chat(messages: List[Dict[str, str]]) -> str:
 
1
  import os
2
  import json
3
+ import subprocess
4
+ import sys
5
  import torch
6
+ from typing import List, Dict
7
+
8
+ # Ensure vllm is installed
9
+ try:
10
+ import vllm
11
+ except ImportError:
12
+ subprocess.check_call([sys.executable, "-m", "pip", "install", "vllm"])
13
+
14
+ # Import the necessary modules after installation
15
  from vllm import LLM, SamplingParams
16
  from vllm.utils import random_uuid
 
17
 
18
  # Function to format chat messages using Qwen's chat template
19
  def format_chat(messages: List[Dict[str, str]]) -> str: