Merge pull request #15 from harishsiravuri/14-update-knowledge-graph

14 update knowledge graph

README.md

@@ -80,3 +80,9 @@ It's time to point you to the official [Documentation Website](https://harishsir
 
 ## Contributing
 If you'd like to contribute, be sure to check out our [contributing guide](./CONTRIBUTING.md)! If you'd like to work on any outstanding items, check out the `roadmap` section in the docs and get started :smiley:
+
+Thanks goes to these incredible people.
+
+<a href="https://github.com/harishsiravuri/kgforge/graphs/contributors">
+  <img style="border-radius: 50%" src="https://contrib.rocks/image?repo=harishsiravuri/kgforge" />
+</a>

kgforge/data_models/data_models.py

@@ -30,6 +30,7 @@ class PromptResponse(BaseModel):
     """
 
     concept: str
+    score: float
     prompt_response: str
 
 

kgforge/kg/kg_construct.py

@@ -1,3 +1,4 @@
+import errno
 import logging
 import os
 from typing import List
@@ -6,6 +7,8 @@
 import networkx as nx
 import transformers
 from transformers import pipeline
+import pickle
+import json
 
 from kgforge.config import KGConfig
 from kgforge.data_models import Prompt, PromptResponse, ResearchArtifact
@@ -125,19 +128,29 @@ def answer_question(
         """
         if artifact is None:
             logger.info("Artifact is needed to answer the question.")
-            return PromptResponse(concept=prompt.concept, prompt_response="Unavailable")
+            return PromptResponse(
+                concept=prompt.concept, score=0, prompt_response="Unavailable"
+            )
         if artifact.full_text is None:
             logger.info("Full text not found.")
-            return PromptResponse(concept=prompt.concept, prompt_response="Unavailable")
+            return PromptResponse(
+                concept=prompt.concept, score=0, prompt_response="Unavailable"
+            )
         if prompt.question == "":
             raise ValueError("Question cannot be empty")
         try:
             nlp = pipeline(task="question-answering", model=self.config.model_name)
             res = nlp(question=prompt.question, context=artifact.full_text)
-            return PromptResponse(concept=prompt.concept, prompt_response=res)
+            return PromptResponse(
+                concept=prompt.concept,
+                score=res.get("score", 0),
+                prompt_response=res.get("answer", "Unavailable"),
+            )
         except transformers.pipelines.base.PipelineException:
             logger.error("Error while answering question")
-            return PromptResponse(concept=prompt.concept, prompt_response="Unavailable")
+            return PromptResponse(
+                concept=prompt.concept, score=0, prompt_response="Unavailable"
+            )
 
     def construct_kg(self) -> None:
         """Constructs knowledge graph using the list of documents
@@ -181,6 +194,69 @@ def construct_kg(self) -> None:
         except Exception as e:
             logger.info("Error while constructing the knowledge graph: " + str(e))
 
+    def read_graph(self, path: str) -> None:
+        """Reads the graph from a file
+
+        Usage example:
+        >>>kg = KnowledgeGraph()
+        >>>kg.read_graph("kg.pickle")
+
+        Args:
+            path (str): Path to the file where the graph is to be read from
+
+        Returns:
+            None: Reads the graph from a file
+
+        Raises:
+            ValueError: If the path is empty
+            FileNotFoundError: If the file is not found
+        """
+        if path is None:
+            raise ValueError("Path cannot be empty")
+        else:
+            if not os.path.isfile(path):
+                raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), path)
+            else:
+                with open(path, "rb") as f:
+                    self.graph = pickle.load(f)
+
+    def write_graph(self, path: str) -> None:
+        """Writes the graph to a file
+
+        Usage example:
+        >>>kg = KnowledgeGraph()
+        >>>kg.write_graph("kg.pickle")
+
+        Args:
+            path (str): Path to the file where the graph is to be written
+
+        Returns:
+            None: Writes the graph to a file
+
+        Raises:
+            ValueError: If the path is empty
+        """
+        try:
+            node_arr = []
+            edge_arr = []
+
+            for node in list(self.graph.nodes(data=True)):
+                node_arr.append(node)
+
+            for edge in list(self.graph.edges()):
+                edge_arr.append(edge)
+
+            graph_dict = {"nodes": node_arr, "edges": edge_arr}
+            with open(path, "w") as f:
+                json.dump(graph_dict, f, indent=4)
+        except:
+            pass
+        # if path is not None and self.graph is not None:
+        #     with open(path, "wb") as f:
+        #         pickle.dump(self.graph, f)
+        # else:
+        #     raise ValueError("Path cannot be empty")
+
     def visualize_kg(self, file_path: str = "graph.png"):
         """Visualizes the knowledge graph
 
@@ -196,6 +272,10 @@ def visualize_kg(self, file_path: str = "graph.png"):
         Raises:
             None
         """
-        pos = nx.spring_layout(self.graph)
-        nx.draw(self.graph, pos=pos, with_labels=True, font_weight="bold")
+        pos = nx.spring_layout(self.graph, k=0.7, iterations=50)
+        nx.draw(self.graph, pos=pos, with_labels=False, font_weight="bold")
+        ax = plt.gca()
+        ax.set_aspect('equal')
+        ax.set_axis_off()
+
         plt.savefig(file_path, format="PNG")