Update app.py

app.py

@@ -249,8 +249,14 @@ class XylariaChat:
 
     def query_knowledge_graph(self, query):
         query_embedding = self.embedding_model.encode(query, convert_to_tensor=True)
-
-        node_embeddings = {node: self.embedding_model.encode(node, convert_to_tensor=True) for node in self.knowledge_graph.nodes()}
+        
+        node_embeddings = {}
+        for node in self.knowledge_graph.nodes():
+            try:
+                node_embedding = self.embedding_model.encode(node, convert_to_tensor=True)
+                node_embeddings[node] = node_embedding
+            except Exception as e:
+                print(f"Error encoding node {node}: {e}")
 
         similarities = {node: util.pytorch_cos_sim(query_embedding, embedding)[0][0].item() for node, embedding in node_embeddings.items()}
 
@@ -340,27 +346,17 @@ class XylariaChat:
             return f"Error during Math OCR: {e}"
 
     def extract_entities_and_relations(self, text):
-        doc = self.embedding_model.tokenizer(text, padding=True, truncation=True, return_tensors="pt")
-
-        with torch.no_grad():
-            outputs = self.embedding_model(**doc)
-
+        inputs = self.embedding_model.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+        
         entities = []
         relations = []
-        for i in range(len(doc['input_ids'][0])):
-            token = self.embedding_model.tokenizer.decode(doc['input_ids'][0][i])
-            if outputs['last_hidden_state'][0][i].norm() > 3:
-                entities.append(token)
 
-        if len(entities) >= 2:
-            for i in range(len(entities) - 1):
-                relation = f"{entities[i]} related_to {entities[i+1]}"
-                relations.append(relation)
+        entities, relations = self.extract_entities_and_relations(message)
+        self.update_knowledge_graph(entities, relations)
 
         return entities, relations
 
-    def update_knowledge_graph(self, text):
-        entities, relations = self.extract_entities_and_relations(text)
+    def update_knowledge_graph(self, entities, relations):
         for entity in entities:
             self.knowledge_graph.add_node(entity)
         for relation in relations:
@@ -372,7 +368,8 @@ class XylariaChat:
 
     def get_response(self, user_input, image=None):
         try:
-            self.update_knowledge_graph(user_input)
+            entities, relations = self.extract_entities_and_relations(user_input)
+            self.update_knowledge_graph(entities, relations)
 
             messages = []