llama-index

1
from llama_index.legacy.extractors.metadata_extractors import (
2
    BaseExtractor,
3
    EntityExtractor,
4
    KeywordExtractor,
5
    QuestionsAnsweredExtractor,
6
    SummaryExtractor,
7
    TitleExtractor,
8
)
9

10

11
def load_extractor(
12
    data: dict,
13
) -> BaseExtractor:
14
    if isinstance(data, BaseExtractor):
15
        return data
16

17
    extractor_name = data.get("class_name", None)
18
    if extractor_name is None:
19
        raise ValueError("Extractor loading requires a class_name")
20

21
    if extractor_name == SummaryExtractor.class_name():
22
        return SummaryExtractor.from_dict(data)
23
    elif extractor_name == QuestionsAnsweredExtractor.class_name():
24
        return QuestionsAnsweredExtractor.from_dict(data)
25
    elif extractor_name == EntityExtractor.class_name():
26
        return EntityExtractor.from_dict(data)
27
    elif extractor_name == TitleExtractor.class_name():
28
        return TitleExtractor.from_dict(data)
29
    elif extractor_name == KeywordExtractor.class_name():
30
        return KeywordExtractor.from_dict(data)
31
    else:
32
        raise ValueError(f"Unknown extractor name: {extractor_name}")
33