| 57 | |
| 58 | |
| 59 | class LLM_Processor: |
| 60 | def __init__(self, args): |
| 61 | self.model = args["llm_model"] |
| 62 | self.base_url = args["llm_url"] |
| 63 | self.api_key = args["llm_api_key"] |
| 64 | self.max_error = args["max_error"] |
| 65 | self.ports = [8001 for i in range(args["gpu_nums"])] # 端口池 |
| 66 | self.gpus = [i for i in range(args["gpu_nums"])] # GPU编号 |
| 67 | if args["use_ollama"]: |
| 68 | self.manager = InstanceManager(self.ports, self.gpus, self.base_url) |
| 69 | self.generate_text = self.manager.generate_text |
| 70 | elif args["use_vllm"]: |
| 71 | self.manager = InstanceManager(self.ports, self.gpus, self.base_url) |
| 72 | self.generate_text = self.vllm_generate_text |
| 73 | else: |
| 74 | self.generate_text = self.default_generate_text |
| 75 | |
| 76 | |
| 77 | def vllm_generate_text(self, prompt, model,max_tokens=4096, output_json=False): |
| 78 | """使用vLLM生成文本""" |
| 79 | |
| 80 | port = self.manager.get_available_instance() |
| 81 | base_url = f"{self.base_url}:{port}/v1" |
| 82 | |
| 83 | try: |
| 84 | if output_json: |
| 85 | # 调用 Chat Completion API 并设置参数 |
| 86 | response = requests.post( |
| 87 | f"{base_url}/chat/completions", |
| 88 | json={ |
| 89 | "model": self.model, |
| 90 | "messages": [ |
| 91 | {"role": "user", "content": prompt} |
| 92 | ], |
| 93 | "max_tokens" : max_tokens, |
| 94 | "response_format":{"type": "json_object"}, |
| 95 | "chat_template_kwargs": {"enable_thinking": False} |
| 96 | }, |
| 97 | timeout=120 |
| 98 | ) |
| 99 | |
| 100 | else: |
| 101 | response = requests.post( |
| 102 | f"{base_url}/chat/completions", |
| 103 | json={ |
| 104 | "model": self.model, |
| 105 | "messages": [ |
| 106 | {"role": "user", "content": prompt} |
| 107 | ], |
| 108 | "max_tokens" : max_tokens, |
| 109 | "chat_template_kwargs": {"enable_thinking": False} |
| 110 | }, |
| 111 | timeout=120 |
| 112 | ) |
| 113 | response.raise_for_status() |
| 114 | res=json.loads(response.content) |
| 115 | response_message = res["choices"][0]["message"]['content'] |
| 116 | |