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}