mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-14 07:59:00 +03:00
40 lines
931 B
Python
40 lines
931 B
Python
from enum import Enum
|
|
|
|
import yaml
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class ParserType(str, Enum):
|
|
"""Parser type enumeration."""
|
|
|
|
UNSTRUCTURED = "unstructured"
|
|
LLAMA_PARSER = "llama_parser"
|
|
MEGAPARSE_VISION = "megaparse_vision"
|
|
|
|
|
|
class StrategyEnum(str, Enum):
|
|
"""Method to use for the conversion"""
|
|
|
|
FAST = "fast"
|
|
AUTO = "auto"
|
|
HI_RES = "hi_res"
|
|
|
|
|
|
class MegaparseBaseConfig(BaseModel):
|
|
@classmethod
|
|
def from_yaml(cls, file_path: str):
|
|
# Load the YAML file
|
|
with open(file_path, "r") as stream:
|
|
config_data = yaml.safe_load(stream)
|
|
|
|
# Instantiate the class using the YAML data
|
|
return cls(**config_data)
|
|
|
|
|
|
class MegaparseConfig(MegaparseBaseConfig):
|
|
method: ParserType = ParserType.UNSTRUCTURED
|
|
strategy: StrategyEnum = StrategyEnum.FAST
|
|
check_table: bool = False
|
|
parsing_instruction: str | None = None
|
|
model_name: str = "gpt-4o"
|