Saving the updated knowledge graph

final.py

@@ -89,7 +89,7 @@ def extract_entities(text, nlp):
     entities = [(ent.text, ent.label_) for ent in doc.ents]
     return entities
 
-def update_knowledge_graph(text, is_real, knowledge_graph, nlp):
+def update_knowledge_graph(text, is_real, knowledge_graph, nlp, save=True, push_to_hf=True):
     """Update knowledge graph with new information"""
     entities = extract_entities(text, nlp)
     for entity, entity_type in entities:
@@ -117,6 +117,16 @@ def update_knowledge_graph(text, is_real, knowledge_graph, nlp):
                 )
             else:
                 knowledge_graph[entity1][entity2]['weight'] += 1
+    if save:
+        from save_model import save_knowledge_graph, push_to_huggingface
+        filepath = save_knowledge_graph(knowledge_graph)
+        
+        # Push to Hugging Face if requested
+        if push_to_hf:
+            repo_id = os.getenv("HF_REPO_ID", "HeheBoi0769/Nexus_NLP_model")
+            push_to_huggingface(filepath, repo_id)
+    
+    return knowledge_graph
 
 def predict_with_knowledge_graph(text, knowledge_graph, nlp):
     """Make predictions using the knowledge graph"""

requirements.txt

@@ -13,4 +13,5 @@ sentencepiece
 timm
 plotly
 networkx
+huggingface_hub
 https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.0/en_core_web_sm-3.7.0.tar.gz

save_model.py

@@ -0,0 +1,43 @@
+import pickle
+import os
+from huggingface_hub import HfApi, login
+import networkx as nx
+
+def save_knowledge_graph(knowledge_graph, filepath="./knowledge_graph_final.pkl"):
+    """Save the knowledge graph to a local file"""
+    # Convert the graph to a serializable format
+    graph_data = {
+        'nodes': {node: data for node, data in knowledge_graph.nodes(data=True)},
+        'edges': {u: {v: data for v, data in knowledge_graph[u].items()} 
+                 for u in knowledge_graph.nodes()}
+    }
+    
+    # Save to file
+    with open(filepath, 'wb') as f:
+        pickle.dump(graph_data, f)
+    
+    print(f"Knowledge graph saved to {filepath}")
+    return filepath
+
+def push_to_huggingface(filepath, repo_id, token=None):
+    """Push the saved knowledge graph to Hugging Face Hub"""
+    if token is None:
+        token = os.getenv("HF_TOKEN")
+        if not token:
+            raise ValueError("No Hugging Face token provided. Set HF_TOKEN environment variable or pass token parameter.")
+    
+    # Login to Hugging Face
+    login(token=token)
+    
+    # Initialize the Hugging Face API
+    api = HfApi()
+    
+    # Upload the file
+    api.upload_file(
+        path_or_fileobj=filepath,
+        path_in_repo="knowledge_graph_final.pkl",
+        repo_id=repo_id,
+        repo_type="model"
+    )
+    
+    print(f"Knowledge graph pushed to Hugging Face Hub: {repo_id}")