From 716e3fe58adee5b8a6bfa91de4b3ba6cf204d172 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Wed, 29 May 2024 00:53:39 +0200
Subject: Wip memory

---
 rag/ui.py | 14 --------------
 1 file changed, 14 deletions(-)

(limited to 'rag/ui.py')

diff --git a/rag/ui.py b/rag/ui.py
index 36e8c4c..ddb3d78 100644
--- a/rag/ui.py
+++ b/rag/ui.py
@@ -13,18 +13,6 @@ from rag.retriever.retriever import Retriever
 from rag.retriever.vector import Document
 
 
-@dataclass
-class Message:
-    role: str
-    message: str
-
-    def as_dict(self, model: str) -> Dict[str, str]:
-        if model == "cohere":
-            return {"role": self.role, "message": self.message}
-        else:
-            return {"role": self.role, "content": self.message}
-
-
 def set_chat_users():
     log.debug("Setting user and bot value")
     ss = st.session_state
@@ -38,13 +26,11 @@ def load_retriever():
     st.session_state.retriever = Retriever()
 
 
-# @st.cache_resource
 def load_generator(model: str):
     log.debug("Loading generator model")
     st.session_state.generator = get_generator(model)
 
 
-# @st.cache_resource
 def load_reranker(model: str):
     log.debug("Loading reranker model")
     st.session_state.reranker = get_reranker(model)
-- 
cgit v1.2.3-70-g09d2