File size: 2,855 Bytes
f2f1a27 285f1d6 84cbb6b f2f1a27 18eefbb f2f1a27 5e316bb f2f1a27 18eefbb f2f1a27 5e316bb 18eefbb f2f1a27 5e316bb 18eefbb f2f1a27 285f1d6 5e316bb 285f1d6 5e316bb 84cbb6b 5e316bb 84cbb6b 5e316bb 84cbb6b 285f1d6 5e316bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import pytest
from marker.converters.pdf import PdfConverter
from marker.services.gemini import GoogleGeminiService
from marker.services.ollama import OllamaService
from marker.services.vertex import GoogleVertexService
from marker.services.openai import OpenAIService
from marker.services.azure_openai import AzureOpenAIService
@pytest.mark.output_format("markdown")
@pytest.mark.config({"page_range": [0]})
def test_empty_llm(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is None
assert pdf_converter.llm_service is None
def test_llm_no_keys(model_dict, config):
with pytest.raises(AssertionError):
PdfConverter(artifact_dict=model_dict, config={"use_llm": True})
@pytest.mark.output_format("markdown")
@pytest.mark.config({"page_range": [0], "use_llm": True, "gemini_api_key": "test"})
def test_llm_gemini(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, GoogleGeminiService)
@pytest.mark.output_format("markdown")
@pytest.mark.config(
{
"page_range": [0],
"use_llm": True,
"vertex_project_id": "test",
"llm_service": "marker.services.vertex.GoogleVertexService",
}
)
def test_llm_vertex(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, GoogleVertexService)
@pytest.mark.output_format("markdown")
@pytest.mark.config(
{
"page_range": [0],
"use_llm": True,
"llm_service": "marker.services.ollama.OllamaService",
}
)
def test_llm_ollama(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, OllamaService)
@pytest.mark.output_format("markdown")
@pytest.mark.config(
{
"page_range": [0],
"use_llm": True,
"llm_service": "marker.services.openai.OpenAIService",
"openai_api_key": "test",
}
)
def test_llm_openai(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, OpenAIService)
@pytest.mark.output_format("markdown")
@pytest.mark.config(
{
"page_range": [0],
"use_llm": True,
"llm_service": "marker.services.azure_openai.AzureOpenAIService",
"azure_endpoint": "https://example.openai.azure.com",
"azure_api_key": "test",
"deployment_name": "test-model",
"azure_api_version": "1",
}
)
def test_llm_azure_openai(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, AzureOpenAIService)
|