File size: 3,324 Bytes
ffb913a 99f5836 ffb913a 99f5836 ffb913a 99f5836 ffb913a 99f5836 ffb913a 99f5836 ffb913a 99f5836 ffb913a 99f5836 ffb913a 99f5836 ffb913a 99f5836 ffb913a 99f5836 ffb913a a25342b ffb913a a25342b 99f5836 ffb913a |
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 |
import os
from lightrag import LightRAG
from lightrag.llm 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",
},
{
"content": "PersonA is a prominent researcher at UniversityB, focusing on artificial intelligence and machine learning.",
"source_id": "Source2",
},
{
"content": "EventY, held in CityC, attracts technology enthusiasts and companies from around the globe.",
"source_id": "Source3",
},
{
"content": "None",
"source_id": "UNKNOWN",
},
],
}
rag.insert_custom_kg(custom_kg)
|