Spaces:
Runtime error
Runtime error
File size: 4,147 Bytes
b699122 |
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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
"""Index registry."""
from typing import Any, Dict, Type
from gpt_index.constants import DATA_KEY, TYPE_KEY
from gpt_index.data_structs.data_structs_v2 import (
KG,
ChatGPTRetrievalPluginIndexDict,
ChromaIndexDict,
CompositeIndex,
EmptyIndex,
FaissIndexDict,
IndexDict,
IndexGraph,
IndexList,
KeywordTable,
MilvusIndexDict,
OpensearchIndexDict,
PineconeIndexDict,
QdrantIndexDict,
SimpleIndexDict,
V2IndexStruct,
WeaviateIndexDict,
)
from gpt_index.data_structs.struct_type import IndexStructType
from gpt_index.data_structs.table_v2 import PandasStructTable, SQLStructTable
from gpt_index.indices.base import BaseGPTIndex, QueryMap
from gpt_index.indices.empty.base import GPTEmptyIndex
from gpt_index.indices.keyword_table.base import GPTKeywordTableIndex
from gpt_index.indices.knowledge_graph.base import GPTKnowledgeGraphIndex
from gpt_index.indices.list.base import GPTListIndex
from gpt_index.indices.struct_store.pandas import GPTPandasIndex
from gpt_index.indices.struct_store.sql import GPTSQLStructStoreIndex
from gpt_index.indices.tree.base import GPTTreeIndex
from gpt_index.indices.vector_store.base import GPTVectorStoreIndex
from gpt_index.indices.vector_store.vector_indices import (
ChatGPTRetrievalPluginIndex,
GPTChromaIndex,
GPTFaissIndex,
GPTMilvusIndex,
GPTOpensearchIndex,
GPTPineconeIndex,
GPTQdrantIndex,
GPTSimpleVectorIndex,
GPTWeaviateIndex,
)
INDEX_STRUCT_TYPE_TO_INDEX_STRUCT_CLASS: Dict[IndexStructType, Type[V2IndexStruct]] = {
IndexStructType.TREE: IndexGraph,
IndexStructType.LIST: IndexList,
IndexStructType.KEYWORD_TABLE: KeywordTable,
IndexStructType.SIMPLE_DICT: SimpleIndexDict,
IndexStructType.DICT: FaissIndexDict,
IndexStructType.WEAVIATE: WeaviateIndexDict,
IndexStructType.PINECONE: PineconeIndexDict,
IndexStructType.QDRANT: QdrantIndexDict,
IndexStructType.MILVUS: MilvusIndexDict,
IndexStructType.CHROMA: ChromaIndexDict,
IndexStructType.OPENSEARCH: OpensearchIndexDict,
IndexStructType.VECTOR_STORE: IndexDict,
IndexStructType.SQL: SQLStructTable,
IndexStructType.PANDAS: PandasStructTable,
IndexStructType.KG: KG,
IndexStructType.EMPTY: EmptyIndex,
IndexStructType.COMPOSITE: CompositeIndex,
IndexStructType.CHATGPT_RETRIEVAL_PLUGIN: ChatGPTRetrievalPluginIndexDict,
}
INDEX_STRUCT_TYPE_TO_INDEX_CLASS: Dict[IndexStructType, Type[BaseGPTIndex]] = {
IndexStructType.TREE: GPTTreeIndex,
IndexStructType.LIST: GPTListIndex,
IndexStructType.KEYWORD_TABLE: GPTKeywordTableIndex,
IndexStructType.SIMPLE_DICT: GPTSimpleVectorIndex,
IndexStructType.DICT: GPTFaissIndex,
IndexStructType.WEAVIATE: GPTWeaviateIndex,
IndexStructType.PINECONE: GPTPineconeIndex,
IndexStructType.QDRANT: GPTQdrantIndex,
IndexStructType.MILVUS: GPTMilvusIndex,
IndexStructType.CHROMA: GPTChromaIndex,
IndexStructType.VECTOR_STORE: GPTVectorStoreIndex,
IndexStructType.SQL: GPTSQLStructStoreIndex,
IndexStructType.PANDAS: GPTPandasIndex,
IndexStructType.KG: GPTKnowledgeGraphIndex,
IndexStructType.EMPTY: GPTEmptyIndex,
IndexStructType.CHATGPT_RETRIEVAL_PLUGIN: ChatGPTRetrievalPluginIndex,
IndexStructType.OPENSEARCH: GPTOpensearchIndex,
}
INDEX_STRUT_TYPE_TO_QUERY_MAP: Dict[IndexStructType, QueryMap] = {
index_type: index_cls.get_query_map()
for index_type, index_cls in INDEX_STRUCT_TYPE_TO_INDEX_CLASS.items()
}
def load_index_struct_from_dict(struct_dict: Dict[str, Any]) -> "V2IndexStruct":
type = struct_dict[TYPE_KEY]
data_dict = struct_dict[DATA_KEY]
cls = INDEX_STRUCT_TYPE_TO_INDEX_STRUCT_CLASS[type]
if type == IndexStructType.COMPOSITE:
struct_dicts: Dict[str, Any] = data_dict["all_index_structs"]
root_id = data_dict["root_id"]
all_index_structs = {
id_: load_index_struct_from_dict(struct_dict)
for id_, struct_dict in struct_dicts.items()
}
return CompositeIndex(all_index_structs=all_index_structs, root_id=root_id)
else:
return cls.from_dict(data_dict)
|