Fu DaiandClaude Opus 4.8
1a585693be
FastAPI service wrapping the Surya-OCR-2 model (datalab-to) served through vLLM: legacy /v1/api/ai/* endpoints, an OpenAI-compatible /v1/chat/completions endpoint, a coalescing request batcher, a local OCR CLI, Docker packaging, multilingual example outputs, and quantization/concurrency benchmarks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
23 lines
784 B
Python
23 lines
784 B
Python
from typing import Optional
|
|
|
|
from transformers import PreTrainedModel
|
|
from transformers.utils import is_flash_attn_2_available
|
|
|
|
|
|
class SuryaPreTrainedModel(PreTrainedModel):
|
|
# No-op if we pass attention, so we can set attention however we want in the config
|
|
def _check_and_adjust_attn_implementation(
|
|
self, attn_implementation: Optional[str], **kwargs
|
|
):
|
|
if attn_implementation is None:
|
|
try:
|
|
self._sdpa_can_dispatch(True)
|
|
attn_implementation = "sdpa"
|
|
except (ValueError, ImportError):
|
|
attn_implementation = "eager"
|
|
|
|
if self._supports_flash_attn and is_flash_attn_2_available():
|
|
attn_implementation = "flash_attention_2"
|
|
|
|
return attn_implementation
|