|
19 | 19 | "import os\n", |
20 | 20 | "\n", |
21 | 21 | "import pandas as pd\n", |
22 | | - "from graphrag.config.enums import ModelType\n", |
23 | 22 | "from graphrag.config.models.drift_search_config import DRIFTSearchConfig\n", |
24 | | - "from graphrag.config.models.language_model_config import LanguageModelConfig\n", |
25 | | - "from graphrag.language_model.manager import ModelManager\n", |
26 | 23 | "from graphrag.query.indexer_adapters import (\n", |
27 | 24 | " read_indexer_entities,\n", |
28 | 25 | " read_indexer_relationships,\n", |
29 | | - " read_indexer_report_embeddings,\n", |
30 | 26 | " read_indexer_reports,\n", |
31 | 27 | " read_indexer_text_units,\n", |
32 | 28 | ")\n", |
|
35 | 31 | ")\n", |
36 | 32 | "from graphrag.query.structured_search.drift_search.search import DRIFTSearch\n", |
37 | 33 | "from graphrag.tokenizer.get_tokenizer import get_tokenizer\n", |
38 | | - "from graphrag_vectors.lancedb import LanceDBVectorStore\n", |
39 | | - "\n", |
40 | | - "INPUT_DIR = \"./inputs/operation dulce\"\n", |
41 | | - "LANCEDB_URI = f\"{INPUT_DIR}/lancedb\"\n", |
42 | | - "\n", |
43 | | - "COMMUNITY_REPORT_TABLE = \"community_reports\"\n", |
44 | | - "COMMUNITY_TABLE = \"communities\"\n", |
45 | | - "ENTITY_TABLE = \"entities\"\n", |
46 | | - "RELATIONSHIP_TABLE = \"relationships\"\n", |
47 | | - "COVARIATE_TABLE = \"covariates\"\n", |
48 | | - "TEXT_UNIT_TABLE = \"text_units\"\n", |
49 | | - "COMMUNITY_LEVEL = 2\n", |
50 | | - "\n", |
51 | | - "\n", |
52 | | - "# read nodes table to get community and degree data\n", |
53 | | - "entity_df = pd.read_parquet(f\"{INPUT_DIR}/{ENTITY_TABLE}.parquet\")\n", |
54 | | - "community_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet\")\n", |
55 | | - "\n", |
56 | | - "print(f\"Entity df columns: {entity_df.columns}\")\n", |
57 | | - "\n", |
58 | | - "entities = read_indexer_entities(entity_df, community_df, COMMUNITY_LEVEL)\n", |
59 | | - "\n", |
60 | | - "# load description embeddings to an in-memory lancedb vectorstore\n", |
61 | | - "# to connect to a remote db, specify url and port values.\n", |
62 | | - "description_embedding_store = LanceDBVectorStore(\n", |
63 | | - " db_uri=LANCEDB_URI,\n", |
64 | | - " index_name=\"entity_description\",\n", |
65 | | - ")\n", |
66 | | - "description_embedding_store.connect()\n", |
67 | | - "\n", |
68 | | - "full_content_embedding_store = LanceDBVectorStore(\n", |
69 | | - " db_uri=LANCEDB_URI,\n", |
70 | | - " index_name=\"community_full_content\",\n", |
71 | | - ")\n", |
72 | | - "full_content_embedding_store.connect()\n", |
73 | | - "\n", |
74 | | - "print(f\"Entity count: {len(entity_df)}\")\n", |
75 | | - "entity_df.head()\n", |
76 | | - "\n", |
77 | | - "relationship_df = pd.read_parquet(f\"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet\")\n", |
78 | | - "relationships = read_indexer_relationships(relationship_df)\n", |
79 | | - "\n", |
80 | | - "print(f\"Relationship count: {len(relationship_df)}\")\n", |
81 | | - "relationship_df.head()\n", |
82 | | - "\n", |
83 | | - "text_unit_df = pd.read_parquet(f\"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet\")\n", |
84 | | - "text_units = read_indexer_text_units(text_unit_df)\n", |
85 | | - "\n", |
86 | | - "print(f\"Text unit records: {len(text_unit_df)}\")\n", |
87 | | - "text_unit_df.head()\n", |
88 | | - "\n", |
89 | | - "report_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet\")\n", |
90 | | - "reports = read_indexer_reports(report_df, community_df, COMMUNITY_LEVEL)\n", |
91 | | - "read_indexer_report_embeddings(reports, full_content_embedding_store)" |
| 34 | + "from graphrag_llm.completion import create_completion\n", |
| 35 | + "from graphrag_llm.config import ModelConfig\n", |
| 36 | + "from graphrag_llm.embedding import create_embedding\n", |
| 37 | + "from graphrag_vectors import IndexSchema, LanceDBVectorStore" |
92 | 38 | ] |
93 | 39 | }, |
94 | 40 | { |
|
99 | 45 | "source": [ |
100 | 46 | "api_key = os.environ[\"GRAPHRAG_API_KEY\"]\n", |
101 | 47 | "\n", |
102 | | - "chat_config = LanguageModelConfig(\n", |
103 | | - " api_key=api_key,\n", |
104 | | - " type=ModelType.Chat,\n", |
| 48 | + "chat_config = ModelConfig(\n", |
| 49 | + " type=\"litellm\",\n", |
105 | 50 | " model_provider=\"openai\",\n", |
106 | 51 | " model=\"gpt-4.1\",\n", |
107 | | - " max_retries=20,\n", |
108 | | - ")\n", |
109 | | - "chat_model = ModelManager().get_or_create_chat_model(\n", |
110 | | - " name=\"local_search\",\n", |
111 | | - " model_type=ModelType.Chat,\n", |
112 | | - " config=chat_config,\n", |
| 52 | + " api_key=api_key,\n", |
113 | 53 | ")\n", |
| 54 | + "chat_model = create_completion(chat_config)\n", |
114 | 55 | "\n", |
115 | 56 | "tokenizer = get_tokenizer(chat_config)\n", |
116 | 57 | "\n", |
117 | | - "embedding_config = LanguageModelConfig(\n", |
118 | | - " api_key=api_key,\n", |
119 | | - " type=ModelType.Embedding,\n", |
| 58 | + "embedding_config = ModelConfig(\n", |
| 59 | + " type=\"litellm\",\n", |
120 | 60 | " model_provider=\"openai\",\n", |
121 | 61 | " model=\"text-embedding-3-large\",\n", |
122 | | - " max_retries=20,\n", |
| 62 | + " api_key=api_key,\n", |
| 63 | + ")\n", |
| 64 | + "\n", |
| 65 | + "text_embedder = create_embedding(embedding_config)" |
| 66 | + ] |
| 67 | + }, |
| 68 | + { |
| 69 | + "cell_type": "code", |
| 70 | + "execution_count": null, |
| 71 | + "metadata": {}, |
| 72 | + "outputs": [], |
| 73 | + "source": [ |
| 74 | + "# parquet files generated from indexing pipeline\n", |
| 75 | + "INPUT_DIR = \"./inputs/operation dulce\"\n", |
| 76 | + "LANCEDB_URI = \"./lancedb\"\n", |
| 77 | + "COMMUNITY_TABLE = \"communities\"\n", |
| 78 | + "COMMUNITY_REPORT_TABLE = \"community_reports\"\n", |
| 79 | + "ENTITY_TABLE = \"entities\"\n", |
| 80 | + "RELATIONSHIP_TABLE = \"relationships\"\n", |
| 81 | + "TEXT_UNIT_TABLE = \"text_units\"\n", |
| 82 | + "COMMUNITY_LEVEL = 2\n", |
| 83 | + "\n", |
| 84 | + "community_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet\")\n", |
| 85 | + "report_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet\")\n", |
| 86 | + "entity_df = pd.read_parquet(f\"{INPUT_DIR}/{ENTITY_TABLE}.parquet\")\n", |
| 87 | + "relationship_df = pd.read_parquet(f\"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet\")\n", |
| 88 | + "text_unit_df = pd.read_parquet(f\"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet\")\n", |
| 89 | + "\n", |
| 90 | + "reports = read_indexer_reports(report_df, community_df, COMMUNITY_LEVEL)\n", |
| 91 | + "entities = read_indexer_entities(entity_df, community_df, COMMUNITY_LEVEL)\n", |
| 92 | + "relationships = read_indexer_relationships(relationship_df)\n", |
| 93 | + "text_units = read_indexer_text_units(text_unit_df)\n", |
| 94 | + "\n", |
| 95 | + "# Connect to the existing GraphRAG vector index for entity-description embeddings.\n", |
| 96 | + "description_embedding_store = LanceDBVectorStore(\n", |
| 97 | + " index_schema=IndexSchema(index_name=\"default-entity-description\")\n", |
123 | 98 | ")\n", |
| 99 | + "description_embedding_store.connect(db_uri=LANCEDB_URI)\n", |
124 | 100 | "\n", |
125 | | - "text_embedder = ModelManager().get_or_create_embedding_model(\n", |
126 | | - " name=\"local_search_embedding\",\n", |
127 | | - " model_type=ModelType.Embedding,\n", |
128 | | - " config=embedding_config,\n", |
| 101 | + "print(\n", |
| 102 | + " f\"Loaded reports={len(reports)}, entities={len(entities)}, relationships={len(relationships)}, text_units={len(text_units)}\"\n", |
129 | 103 | ")" |
130 | 104 | ] |
131 | 105 | }, |
|
188 | 162 | ], |
189 | 163 | "metadata": { |
190 | 164 | "kernelspec": { |
191 | | - "display_name": "Python 3", |
| 165 | + "display_name": "graphrag-monorepo (3.12.10)", |
192 | 166 | "language": "python", |
193 | 167 | "name": "python3" |
194 | 168 | }, |
|
0 commit comments