58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
import os
|
|
import dspy
|
|
from dotenv import load_dotenv
|
|
from modaic import PrecompiledProgram, PrecompiledConfig
|
|
from src.neo4j import Neo4j
|
|
|
|
load_dotenv()
|
|
|
|
# set up Neo4j using NEO4J_URI
|
|
neo4j = Neo4j(
|
|
uri=os.getenv("NEO4J_URI"),
|
|
user=os.getenv("NEO4J_USER"),
|
|
password=os.getenv("NEO4J_PASSWORD"),
|
|
)
|
|
|
|
class CypherFromText(dspy.Signature):
|
|
"""Instructions:
|
|
Create a Cypher MERGE statement to model all entities and relationships found in the text following these guidelines:
|
|
- Refer to the provided schema and use existing or similar nodes, properties or relationships before creating new ones.
|
|
- Use generic categories for node and relationship labels."""
|
|
text = dspy.InputField(
|
|
desc="Text to model using nodes, properties and relationships."
|
|
)
|
|
neo4j_schema = dspy.InputField(
|
|
desc="Current graph schema in Neo4j as a list of NODES and RELATIONSHIPS."
|
|
)
|
|
statement = dspy.OutputField(
|
|
desc="Cypher statement to merge nodes and relationships found in the text."
|
|
)
|
|
|
|
class GenerateCypherConfig(PrecompiledConfig):
|
|
neo4j_schema: list[str] = []
|
|
model: str = "openai/gpt-4o"
|
|
max_tokens: int = 1024
|
|
|
|
|
|
class GenerateCypher(PrecompiledProgram):
|
|
config: GenerateCypherConfig
|
|
|
|
def _init_(self, config: GenerateCypherConfig, **kwargs):
|
|
super()._init_(**kwargs)
|
|
self.lm = dspy.LM(
|
|
model=config.model,
|
|
max_tokens=config.max_tokens,
|
|
)
|
|
self.generate_cypher = dspy.ChainOfThought(CypherFromText)
|
|
self.generate_cypher.set_lm(self.lm)
|
|
|
|
def forward(self, text: str, neo4j_schema: list[str]):
|
|
return self.generate_cypher(text=text, neo4j_schema=neo4j_schema)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
generate_cypher = GenerateCypher(GenerateCypherConfig())
|
|
generate_cypher.push_to_hub("farouk1/text-to-cypher", with_code=True, tag="v0.0.3", commit_message="set LM")
|
|
|
|
|