parser.py support command parameter: protocol

This commit is contained in:
l.tingting
2025-09-02 11:25:06 +08:00
parent 2562ed0df4
commit ee5da7d955
2 changed files with 10 additions and 1 deletions
+2 -1
View File
@@ -12,6 +12,7 @@ import os
def inference_with_vllm( def inference_with_vllm(
image, image,
prompt, prompt,
protocol="http",
ip="localhost", ip="localhost",
port=8000, port=8000,
temperature=0.1, temperature=0.1,
@@ -20,7 +21,7 @@ def inference_with_vllm(
model_name='model', model_name='model',
): ):
addr = f"http://{ip}:{port}/v1" addr = f"{protocol}://{ip}:{port}/v1"
client = OpenAI(api_key="{}".format(os.environ.get("API_KEY", "0")), base_url=addr) client = OpenAI(api_key="{}".format(os.environ.get("API_KEY", "0")), base_url=addr)
messages = [] messages = []
messages.append( messages.append(
+8
View File
@@ -20,6 +20,7 @@ class DotsOCRParser:
""" """
def __init__(self, def __init__(self,
protocol='http',
ip='localhost', ip='localhost',
port=8000, port=8000,
model_name='model', model_name='model',
@@ -36,6 +37,7 @@ class DotsOCRParser:
self.dpi = dpi self.dpi = dpi
# default args for vllm server # default args for vllm server
self.protocol = protocol
self.ip = ip self.ip = ip
self.port = port self.port = port
self.model_name = model_name self.model_name = model_name
@@ -119,6 +121,7 @@ class DotsOCRParser:
image, image,
prompt, prompt,
model_name=self.model_name, model_name=self.model_name,
protocol=self.protocol,
ip=self.ip, ip=self.ip,
port=self.port, port=self.port,
temperature=self.temperature, temperature=self.temperature,
@@ -347,6 +350,10 @@ def main():
metavar=('x1', 'y1', 'x2', 'y2'), metavar=('x1', 'y1', 'x2', 'y2'),
help='should give this argument if you want to prompt_grounding_ocr' help='should give this argument if you want to prompt_grounding_ocr'
) )
parser.add_argument(
"--protocol", type=str, choices=['http', 'https'], default="http",
help=""
)
parser.add_argument( parser.add_argument(
"--ip", type=str, default="localhost", "--ip", type=str, default="localhost",
help="" help=""
@@ -398,6 +405,7 @@ def main():
args = parser.parse_args() args = parser.parse_args()
dots_ocr_parser = DotsOCRParser( dots_ocr_parser = DotsOCRParser(
protocol=args.protocol,
ip=args.ip, ip=args.ip,
port=args.port, port=args.port,
model_name=args.model_name, model_name=args.model_name,