Update ccat_reranker.py

nickprock · Jun 2, 2024 · 00eae1a · 00eae1a
1 parent 2ca237c
commit 00eae1a
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/ccat_reranker.py b/ccat_reranker.py
@@ -23,7 +23,8 @@ def after_cat_recalls_memories(cat) -> None:
             cat.working_memory['episodic_memories'] = recent_docs
         else:
             print("#HicSuntGattones")
-
+
+    #TODO refactor
     if settings["SBERT"]:
         model = CrossEncoder(settings["ranker"])
         if cat.working_memory['declarative_memories']:
@@ -35,6 +36,15 @@ def after_cat_recalls_memories(cat) -> None:
                 cat.working_memory['declarative_memories'] = sbert_docs
         else:
             print("#HicSuntGattones")
+    else:
+        if cat.working_memory['declarative_memories']:
+            if settings["LITM"]:
+                litm_docs = litm(cat.working_memory['declarative_memories'])
+                cat.working_memory['declarative_memories'] = litm_docs
+            else:
+                print("#HicSuntGattones")
+        else:
+            print("#HicSuntGattones")
 
     if settings["FILTER"]:
         if cat.working_memory['procedural_memories']: