Update ocr_vllm_server.py

magic_pdf/dict2md/ocr_vllm_server.py

@@ -14,7 +14,8 @@ from transformers import AutoProcessor
 from vllm import LLM, SamplingParams
 import os
 import configparser
-
+import time
+from loguru import logger
 from fastapi import FastAPI
 from pydantic import BaseModel
 from typing import Optional
@@ -281,13 +282,14 @@ async def predict(item: Item):
             ]
         }
     ]
-
+    start = time.time()
     generated_text = ''
     for response in call_local_model(llm, processor, messages):
         generated_text = unparse_text(response)
 
     _gc()
-
+    end = time.time()
+    logger.info(f'【{messages}】\n 解析的结果是：{generated_text},耗时为：{end-start}')
     return {"Generated Text": generated_text}
 
 
@@ -301,4 +303,4 @@ if __name__ == "__main__":
     host, port = config.get('server', 'ocr_server').split('://')[1].split(':')[0], int(
         config.get('server', 'ocr_server').split('://')[1].split(':')[1])
     # port = int(config.get('server', 'ocr_port'))
-    uvicorn.run(app, host=host, port=port)
\ No newline at end of file
+    uvicorn.run(app, host=host, port=port)