chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,115 @@
|
||||
import os
|
||||
from lightrag import LightRAG
|
||||
from lightrag.llm.openai import gpt_4o_mini_complete
|
||||
#########
|
||||
# Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert()
|
||||
# import nest_asyncio
|
||||
# nest_asyncio.apply()
|
||||
#########
|
||||
|
||||
WORKING_DIR = "./custom_kg"
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model
|
||||
# llm_model_func=gpt_4o_complete # Optionally, use a stronger model
|
||||
)
|
||||
|
||||
custom_kg = {
|
||||
"entities": [
|
||||
{
|
||||
"entity_name": "CompanyA",
|
||||
"entity_type": "Organization",
|
||||
"description": "A major technology company",
|
||||
"source_id": "Source1",
|
||||
},
|
||||
{
|
||||
"entity_name": "ProductX",
|
||||
"entity_type": "Product",
|
||||
"description": "A popular product developed by CompanyA",
|
||||
"source_id": "Source1",
|
||||
},
|
||||
{
|
||||
"entity_name": "PersonA",
|
||||
"entity_type": "Person",
|
||||
"description": "A renowned researcher in AI",
|
||||
"source_id": "Source2",
|
||||
},
|
||||
{
|
||||
"entity_name": "UniversityB",
|
||||
"entity_type": "Organization",
|
||||
"description": "A leading university specializing in technology and sciences",
|
||||
"source_id": "Source2",
|
||||
},
|
||||
{
|
||||
"entity_name": "CityC",
|
||||
"entity_type": "Location",
|
||||
"description": "A large metropolitan city known for its culture and economy",
|
||||
"source_id": "Source3",
|
||||
},
|
||||
{
|
||||
"entity_name": "EventY",
|
||||
"entity_type": "Event",
|
||||
"description": "An annual technology conference held in CityC",
|
||||
"source_id": "Source3",
|
||||
},
|
||||
],
|
||||
"relationships": [
|
||||
{
|
||||
"src_id": "CompanyA",
|
||||
"tgt_id": "ProductX",
|
||||
"description": "CompanyA develops ProductX",
|
||||
"keywords": "develop, produce",
|
||||
"weight": 1.0,
|
||||
"source_id": "Source1",
|
||||
},
|
||||
{
|
||||
"src_id": "PersonA",
|
||||
"tgt_id": "UniversityB",
|
||||
"description": "PersonA works at UniversityB",
|
||||
"keywords": "employment, affiliation",
|
||||
"weight": 0.9,
|
||||
"source_id": "Source2",
|
||||
},
|
||||
{
|
||||
"src_id": "CityC",
|
||||
"tgt_id": "EventY",
|
||||
"description": "EventY is hosted in CityC",
|
||||
"keywords": "host, location",
|
||||
"weight": 0.8,
|
||||
"source_id": "Source3",
|
||||
},
|
||||
],
|
||||
"chunks": [
|
||||
{
|
||||
"content": "ProductX, developed by CompanyA, has revolutionized the market with its cutting-edge features.",
|
||||
"source_id": "Source1",
|
||||
"source_chunk_index": 0,
|
||||
},
|
||||
{
|
||||
"content": "One outstanding feature of ProductX is its advanced AI capabilities.",
|
||||
"source_id": "Source1",
|
||||
"chunk_order_index": 1,
|
||||
},
|
||||
{
|
||||
"content": "PersonA is a prominent researcher at UniversityB, focusing on artificial intelligence and machine learning.",
|
||||
"source_id": "Source2",
|
||||
"source_chunk_index": 0,
|
||||
},
|
||||
{
|
||||
"content": "EventY, held in CityC, attracts technology enthusiasts and companies from around the globe.",
|
||||
"source_id": "Source3",
|
||||
"source_chunk_index": 0,
|
||||
},
|
||||
{
|
||||
"content": "None",
|
||||
"source_id": "UNKNOWN",
|
||||
"source_chunk_index": 0,
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
rag.insert_custom_kg(custom_kg)
|
||||
Reference in New Issue
Block a user