wehub-resource-sync
182 lines
6.6 KiB
Python
182 lines
6.6 KiB
Python
"""
|
|
DeepSeek OCR Parser
|
|
|
|
Uses DeepSeek-OCR via SiliconFlow API for document parsing and OCR.
|
|
"""
|
|
|
|
import base64
|
|
import os
|
|
import re
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import fitz # PyMuPDF
|
|
import requests
|
|
|
|
from yuxi.knowledge.parser.base import BaseDocumentProcessor, DocumentParserException
|
|
from yuxi.utils import logger
|
|
|
|
|
|
class DeepSeekOCRParser(BaseDocumentProcessor):
|
|
"""DeepSeek OCR Parser using SiliconFlow API"""
|
|
|
|
# MIME type mapping for supported formats
|
|
MIME_TYPE_MAP = {
|
|
".pdf": "application/pdf",
|
|
".png": "image/png",
|
|
".jpg": "image/jpeg",
|
|
".jpeg": "image/jpeg",
|
|
".bmp": "image/bmp",
|
|
".webp": "image/webp",
|
|
}
|
|
|
|
def __init__(self, api_key: str | None = None):
|
|
self.api_key = api_key or os.getenv("SILICONFLOW_API_KEY")
|
|
if not self.api_key:
|
|
raise DocumentParserException(
|
|
"SILICONFLOW_API_KEY environment variable not set", "deepseek_ocr", "missing_api_key"
|
|
)
|
|
|
|
self.api_url = "https://api.siliconflow.cn/v1/chat/completions"
|
|
self.model = "deepseek-ai/DeepSeek-OCR"
|
|
|
|
self.headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
}
|
|
|
|
def get_service_name(self) -> str:
|
|
return "deepseek_ocr"
|
|
|
|
def get_supported_extensions(self) -> list[str]:
|
|
"""DeepSeek OCR supports PDF and images"""
|
|
return list(self.MIME_TYPE_MAP.keys())
|
|
|
|
def check_health(self) -> dict[str, Any]:
|
|
"""Check API availability and key validity"""
|
|
try:
|
|
# We can't easily "ping" without cost, but we can check if the model list is accessible
|
|
models_url = "https://api.siliconflow.cn/v1/models"
|
|
response = requests.get(models_url, headers=self.headers, timeout=10)
|
|
|
|
if response.status_code == 200:
|
|
return {
|
|
"status": "healthy",
|
|
"message": "DeepSeek OCR (SiliconFlow) is available",
|
|
"details": {"api_url": self.api_url},
|
|
}
|
|
elif response.status_code == 401:
|
|
return {"status": "unhealthy", "message": "Invalid API Key", "details": {"error_code": "401"}}
|
|
else:
|
|
return {
|
|
"status": "unhealthy",
|
|
"message": f"API Error: {response.status_code}",
|
|
"details": {"status_code": response.status_code},
|
|
}
|
|
except Exception as e:
|
|
return {"status": "unavailable", "message": f"Connection failed: {str(e)}", "details": {"error": str(e)}}
|
|
|
|
def process_file(self, file_path: str, params: dict[str, Any] | None = None) -> str:
|
|
"""
|
|
Process file using DeepSeek OCR via SiliconFlow
|
|
"""
|
|
if not os.path.exists(file_path):
|
|
raise DocumentParserException(f"File not found: {file_path}", self.get_service_name(), "file_not_found")
|
|
|
|
file_ext = Path(file_path).suffix.lower()
|
|
if not self.supports_file_type(file_ext):
|
|
raise DocumentParserException(
|
|
f"Unsupported file type: {file_ext}", self.get_service_name(), "unsupported_file_type"
|
|
)
|
|
|
|
try:
|
|
start_time = time.time()
|
|
logger.info(f"DeepSeek OCR starting: {os.path.basename(file_path)}")
|
|
|
|
if file_ext == ".pdf":
|
|
content = self._process_pdf(file_path)
|
|
else:
|
|
content = self._process_image(file_path)
|
|
|
|
processing_time = time.time() - start_time
|
|
logger.info(
|
|
f"DeepSeek OCR finished: {os.path.basename(file_path)} - {len(content)} chars ({processing_time:.2f}s)"
|
|
)
|
|
|
|
return content
|
|
|
|
except Exception as e:
|
|
if isinstance(e, DocumentParserException):
|
|
raise
|
|
error_msg = f"DeepSeek OCR failed: {str(e)}"
|
|
logger.error(error_msg)
|
|
raise DocumentParserException(error_msg, self.get_service_name(), "processing_failed")
|
|
|
|
def _process_pdf(self, file_path: str) -> str:
|
|
"""Process PDF by converting pages to images"""
|
|
doc = fitz.open(file_path)
|
|
try:
|
|
full_text = []
|
|
|
|
total_pages = len(doc)
|
|
logger.info(f"Processing PDF with {total_pages} pages")
|
|
|
|
for i, page in enumerate(doc):
|
|
logger.debug(f"Processing page {i + 1}/{total_pages}")
|
|
# Convert page to image (200 DPI for better quality)
|
|
pix = page.get_pixmap(dpi=200)
|
|
img_bytes = pix.tobytes("png")
|
|
|
|
page_text = self._call_api(img_bytes, "image/png")
|
|
full_text.append(page_text)
|
|
|
|
return "\n\n".join(full_text)
|
|
finally:
|
|
doc.close()
|
|
|
|
def _process_image(self, file_path: str) -> str:
|
|
"""Process single image file"""
|
|
mime_type = self._get_mime_type(file_path)
|
|
with open(file_path, "rb") as f:
|
|
file_content = f.read()
|
|
return self._call_api(file_content, mime_type)
|
|
|
|
def _call_api(self, data_bytes: bytes, mime_type: str) -> str:
|
|
"""Call SiliconFlow API"""
|
|
encoded_string = base64.b64encode(data_bytes).decode("utf-8")
|
|
data_url = f"data:{mime_type};base64,{encoded_string}"
|
|
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image_url", "image_url": {"url": data_url}},
|
|
{"type": "text", "text": "<image>\n<|grounding|>Convert the document to markdown. "},
|
|
],
|
|
}
|
|
]
|
|
|
|
payload = {"model": self.model, "messages": messages, "max_tokens": 4096, "temperature": 0.1}
|
|
|
|
response = requests.post(self.api_url, headers=self.headers, json=payload, timeout=120)
|
|
|
|
if response.status_code != 200:
|
|
error_msg = f"API Error {response.status_code}: {response.text}"
|
|
logger.error(error_msg)
|
|
raise DocumentParserException(error_msg, self.get_service_name(), f"http_{response.status_code}")
|
|
|
|
result = response.json()
|
|
content = result["choices"][0]["message"]["content"]
|
|
|
|
# Clean up special tags like <|ref|>...<|/ref|> and <|det|>...<|/det|>
|
|
content = re.sub(r"<\|ref\|>.*?<\|/ref\|>", "", content)
|
|
content = re.sub(r"<\|det\|>.*?<\|/det\|>", "", content)
|
|
# content = re.sub(r"<\|.*?\|>", "", content)
|
|
|
|
return content.strip()
|
|
|
|
def _get_mime_type(self, file_path: str) -> str:
|
|
file_ext = Path(file_path).suffix.lower()
|
|
return self.MIME_TYPE_MAP.get(file_ext, "image/jpeg") # Default fallback
|