大模型流式推理

目录http 流式推理直接返回http 流式推理ocr_stream_client.py# codingutf-8 import base64 import json import os import urllib.request import logging import logging.handlers from pathlib import Path from typing import Optional, List, Dict, Any, AsyncGenerator import requests LOG_DIR Path(__file__).resolve().parent / logs LOG_DIR.mkdir(exist_okTrue) base_url http://192.168.100.203:17890 # 统一使用第二个地址 # 主日志处理器geogebra_api.log main_handler logging.handlers.TimedRotatingFileHandler( filenameLOG_DIR / geogebra_api.log, whenmidnight, interval1, backupCount30, encodingutf-8 ) main_handler.setFormatter(logging.Formatter(%(asctime)s - %(name)s - %(levelname)s - %(message)s)) stream_handler logging.StreamHandler() stream_handler.setFormatter(logging.Formatter(%(asctime)s - %(name)s - %(levelname)s - %(message)s)) logging.basicConfig(levellogging.INFO, handlers[stream_handler, main_handler]) logger logging.getLogger(__name__) _MODEL_ID None def get_model_id() - str: 获取 vLLM 服务的模型 ID并缓存 global _MODEL_ID if _MODEL_ID is not None: return _MODEL_ID try: resp urllib.request.urlopen(f{base_url}/v1/models, timeout30) data json.loads(resp.read().decode(utf-8)) _MODEL_ID data[data][0][id] logger.info(f获取到模型 ID: {_MODEL_ID}) return _MODEL_ID except Exception as e: logger.error(f获取模型 ID 失败: {e}) # raise HTTPException(status_code500, detail无法获取 vLLM 模型 ID) model_id get_model_id() def _http_stream(url: str, payload: dict, timeout: int 300): 流式请求vLLM接口逐行返回SSE数据 data json.dumps(payload).encode(utf-8) req urllib.request.Request(url, datadata, headers{Content-Type: application/json}) with urllib.request.urlopen(req, timeouttimeout) as resp: for line in resp: line line.decode(utf-8).strip() if not line: continue if line.startswith(data: ): data_str line[6:] # 去掉 data: 前缀 if data_str [DONE]: break try: chunk json.loads(data_str) yield chunk except json.JSONDecodeError: continue prompt_file ocr_prompt.md with open(prompt_file, r, encodingutf-8) as f: system_prompt f.read() # 一行搞定包含所有换行 def encode_image_to_base64(image_path: str) - str: 将图片编码为 base64 with open(image_path, rb) as image_file: return base64.b64encode(image_file.read()).decode(utf-8) def ocr_image_processing(img_path: str, model_id: str, base_url: str, session_id: str, rep_id: str, timeout: int 300) - AsyncGenerator[str, None]: try: # 读取图片数据 with open(img_path, rb) as f: image_data f.read() # 获取图片文件名 image_filename os.path.basename(img_path) json_log_path image_filename[:-4] .json # 构建消息 ocr_messages [{role: user, content: [{type: text, text: 请识别图片中的文字和数学公式}, {type: image_url, image_url: {url: fdata:image/jpeg;base64,{base64.b64encode(image_data).decode(utf-8)}}}]}] full_ocr_messages [{role: system, content: system_prompt}] ocr_messages # 构建请求负载 ocr_payload {model: model_id, messages: full_ocr_messages, temperature: 0.1, max_tokens: 2048, stream: True} # 直接使用 requests 发送请求 response requests.post(f{base_url}/v1/chat/completions, jsonocr_payload, streamTrue, timeouttimeout) if response.status_code ! 200: error_msg fHTTP错误: {response.status_code} logger.error(error_msg) yield fdata: {json.dumps({type: ocr_error, error: error_msg, session_id: session_id, rep_id: rep_id},ensure_asciiFalse)}\n\n return ocr_full_content ocr_success False # 处理流式响应 for line in response.iter_lines(): if not line: continue line line.decode(utf-8).strip() if not line.startswith(data: ): continue data_str line[6:] # 去掉 data: 前缀 if data_str [DONE]: break try: chunk json.loads(data_str) choices chunk.get(choices, []) if not choices: continue delta choices[0].get(delta, {}) content delta.get(content, ) if content: ocr_full_content content print(content) yield fdata: {json.dumps({type: ocr_chunk, content: content, session_id: session_id, rep_id: rep_id},ensure_asciiFalse)}\n\n except json.JSONDecodeError: continue ocr_text ocr_full_content.strip() if ocr_text: ocr_success True logger.info(fOCR识别成功长度: {len(ocr_text)}) # 保存记录 record {session_id: session_id, image_filename: image_filename, ocr_text: ocr_text, success: ocr_success} with open(json_log_path, w, encodingutf-8) as f: json.dump(record, f, ensure_asciiFalse, indent2) logger.info(fOCR记录已保存: {json_log_path}) logger.info(fOCR识别文本内容:\n{ocr_text}) yield fdata: {json.dumps({type: ocr_done, content: ocr_text, session_id: session_id, rep_id: rep_id},ensure_asciiFalse)}\n\n else: logger.warning(OCR识别结果为空) yield fdata: {json.dumps({type: ocr_error, error: OCR识别结果为空, session_id: session_id, rep_id: rep_id},ensure_asciiFalse)}\n\n except FileNotFoundError: error_msg f图片文件不存在: {img_path} logger.error(error_msg) yield fdata: {json.dumps({type: ocr_error, error: error_msg, session_id: session_id, rep_id: rep_id},ensure_asciiFalse)}\n\n except Exception as e: error_msg fOCR处理异常: {str(e)} logger.error(error_msg) yield fdata: {json.dumps({type: ocr_error, error: error_msg, session_id: session_id, rep_id: rep_id},ensure_asciiFalse)}\n\n if __name__ __main__: # 配置日志 logging.basicConfig(levellogging.INFO) # 配置参数 img_path rC:\Users\ChanJing-01\Documents\math\07\hanshu.png session_id session_id rep_id rep_id # 模拟的http_stream_function实际使用时需要替换为真实的请求函数 def mock_http_stream_function(url, payload, timeout): # 模拟返回数据 yield {choices: [{delta: {content: 这是一个测试}}]} yield {choices: [{delta: {content: OCR识别结果}}]} # 调用函数 # for chunk in ocr_image_processing(img_pathimg_path, model_idmodel_id, base_urlbase_url, session_idsession_id, rep_idrep_id, http_stream_functionmock_http_stream_function): for chunk in ocr_image_processing(img_pathimg_path, model_idmodel_id, base_urlbase_url, session_idsession_id, rep_idrep_id): print(chunk)直接返回ocr_from_server.pyimport urllib from http.client import HTTPException from pathlib import Path import requests import json import base64 from typing import List, Dict, Any import time def encode_image_to_base64(image_path: str) - str: 将图片编码为 base64 with open(image_path, rb) as image_file: return base64.b64encode(image_file.read()).decode(utf-8) _MODEL_ID None def get_model_id() - str: 获取 vLLM 服务的模型 ID并缓存 global _MODEL_ID if _MODEL_ID is not None: return _MODEL_ID try: resp urllib.request.urlopen(f{BASE_URL}/v1/models, timeout30) data json.loads(resp.read().decode(utf-8)) _MODEL_ID data[data][0][id] print(f获取到模型 ID: {_MODEL_ID}) return _MODEL_ID except Exception as e: print(f获取模型 ID 失败: {e}) raise HTTPException(status_code500, detail无法获取 vLLM 模型 ID) def chat_with_images(text: str, image_paths: List[str], max_tokens: int 1024, temperature: float 0.7) - str: model_id get_model_id() # 构建 content 列表 content [] # 1. 添加文本内容 content.append({type: text, text: text}) for image_path in image_paths: # 如果是本地图片转换为 base64 if image_path.startswith((http://, https://)): # URL 图片 content.append({type: image_url, image_url: {url: image_path}}) else: # 本地图片 base64_image encode_image_to_base64(image_path) content.append({type: image_url, image_url: {url: fdata:image/jpeg;base64,{base64_image}}}) prompt_file ocr_prompt.md with open(prompt_file, r, encodingutf-8) as f: system_prompt f.read() # 一行搞定包含所有换行 payload { model: model_id, messages: [ {role: system, content: system_prompt}, {role: user, content: content}, ], temperature: 0.0, max_tokens: 4096, stream: True, } print(f Sending request with {len(image_paths)} images...) print(f Text: {text}) print(- * 80) full_response try: response requests.post(f{BASE_URL}/v1/chat/completions, jsonpayload, streamTrue, timeout300) if response.status_code ! 200: print(f❌ Error: {response.status_code}) print(response.text) return # 5. 处理流式响应 for line in response.iter_lines(): if line: line line.decode(utf-8) if line.startswith(data: ): data line[6:] # 去掉 data: 前缀 if data [DONE]: break try: chunk json.loads(data) choices chunk.get(choices, []) if choices: delta choices[0].get(delta, {}) content_delta delta.get(content, ) if content_delta: print(content_delta, end, flushTrue) full_response content_delta except json.JSONDecodeError: continue print(\n - * 80) return full_response except requests.exceptions.Timeout: print(❌ Request timeout) return except Exception as e: print(f❌ Error: {e}) return if __name__ __main__: BASE_URL http://192.168.100.203:17890 image_paths [rC:\Users\ChanJing-01\Documents\math\07\a12.jpg] image_paths [rC:\Users\ChanJing-01\Documents\math\07\hanshu.png] # response chat_with_images(text图中有两个矩形把几何图复现一下, image_paths[rC:\Users\ChanJing-01\Pictures\wav_sucai\ca1e63ed-1560-4ade-9e45-57aa500b7123.png]) response chat_with_images(text请识别图片中的文字和数学公式, image_pathsimage_paths) print(response)