(self, args)
| 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): |
nothing calls this directly
no test coverage detected