refonte du mécanisme mnésique + élision des listes pr dict

Author: stadiello

.gitignore

@@ -0,0 +1 @@
+venv/

src/shortterm_memory/ChatbotMemory.py

@@ -1,20 +1,23 @@
 import torch
-from transformers import pipeline, BartTokenizer, BartForConditionalGeneration
+from transformers import BartTokenizer, BartForConditionalGeneration
 import logging
 
-# Set device to GPU if available
-if torch.cuda.is_available():
-    device:int = 0
-else:
-    device:int = -1
+# Configuration du logging
+logging.basicConfig(level=logging.INFO)
 
-MAX_MEMORY_SIZE:int = 2000
+# Détection automatique du device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Paramètres globaux
+MAX_MEMORY_SIZE = 2000  # Limite du nombre de messages
+MAX_TOKENS_PER_MESSAGE = 1000  # Limite pour compresser la mémoire
+BATCH_SIZE = 5  # Taille du batch pour la compression
 
 class ChatbotMemory:
-    def __init__(self, conv:list = []):
-        self.conversation_history = conv
+    def __init__(self, conv:list=None):
+        self.conversation_history = conv or []
         self.tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')
-        self.model = BartForConditionalGeneration.from_pretrained('facebook/bart-large-cnn')
+        self.model = BartForConditionalGeneration.from_pretrained('facebook/bart-large-cnn').to(device)
 
     def update_memory(self, user_input:str, bot_response:str)->None:
         """
@@ -23,18 +26,16 @@ def update_memory(self, user_input:str, bot_response:str)->None:
             user_input (str): The input provided by the user.
             bot_response (str): The response generated by the Chatbot.
         Returns:
-            None
-        """
-        self.conversation_history.append(f"'user': {user_input}, 'bot': {bot_response}")
-        
-        if self.memory_counter(self.conversation_history) > 1000:
-            self.conversation_history = self.compressed_memory(self.conversation_history)
-            logging.info("Memory compressed.")
-        
+            None        """
+        self.conversation_history.append({'user': user_input, 'bot': bot_response})
+
+        if self.memory_counter() > MAX_TOKENS_PER_MESSAGE:
+            self.conversation_history = self.compressed_memory()
+            logging.info("Mémoire compressée.")
+
         if len(self.conversation_history) > MAX_MEMORY_SIZE:
             self.conversation_history.pop(0)
-            logging.info("Memory trimmed.")  
-        return 0
+            logging.info("Mémoire tronquée.")
 
     def get_memory(self):
         """
@@ -44,105 +45,51 @@ def get_memory(self):
             The conversation history.
         """
         return self.conversation_history
-    
-    def _get_compressed_memory(self, text):
-        
-        inputs = self.tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=1024, truncation=True)
-        summary_ids = self.model.generate(inputs, max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
-        summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-        return summary
-# def _get_compressed_memory(sentence:str)->str:
-#     """
-#     Compresses the input sentence using the Facebook BART model for summarization.
-
-#     Args:
-#         sentence: The input sentence to be compressed.
 
-#     Returns:
-#         str: The compressed summary of the input sentence.
-#     """
-#     summarizer:str = pipeline("summarization",model="facebook/bart-large-cnn",device=device)
-#     summary:str = summarizer(sentence, max_length=50, min_length=5, do_sample=False)
-#     return summary[0]['summary_text']
-
-    def compressed_memory(self, conv_hist:list)->list:
+    def _get_compressed_memory(self, text:str):
         """
-        Compresses each sentence in the conversation history list using summarization.
-
-        Args:
-            conv_hist: List of sentences representing the conversation history.
-
-        Returns:
-            list: List of compressed summaries for each sentence in the conversation history.
+        Résume un bloc de texte.
         """
-        # return [_get_compressed_memory(sentence) for sentence in conv_hist]
-        return [self._get_compressed_memory(' '.join(conv_hist[i:i+5])) for i in range(0, len(conv_hist), 5)]
-
-
-    def memory_counter(self, conv_hist:list[str])->int:
+        inputs = self.tokenizer(
+            f"summarize: {text}",
+            return_tensors="pt",
+            max_length=1024,
+            truncation=True,
+        ).to(device)
+
+        summary_ids = self.model.generate(
+            inputs.input_ids,
+            max_length=150,
+            min_length=40,
+            length_penalty=2.0,
+            num_beams=4,
+            early_stopping=True
+        )
+        return self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+
+    def compressed_memory(self):
         """
-        Counts the total number of words in the conversation history list.
-
-        Args:
-            conv_hist: List of sentences representing the conversation history.
+        Résume l'historique de la conversation par batch de BATCH_SIZE.
+        """
+        combined_history = [f"User: {entry['user']} Bot: {entry['bot']}" for entry in self.conversation_history]
+        return [
+            {'summary': self._get_compressed_memory(' '.join(combined_history[i:i+BATCH_SIZE]))}
+            for i in range(0, len(combined_history), BATCH_SIZE)
+        ]
 
-        Returns:
-            int: Total number of words in the conversation history.
+    def memory_counter(self):
+        """
+        Compte le nombre total de tokens dans l'historique.
         """
-        st = ''.join(conv_hist)
-        return len(st.split())
+        all_text = ' '.join(f"{entry['user']} {entry['bot']}" for entry in self.conversation_history)
+        return len(self.tokenizer.tokenize(all_text))
 
 if __name__ == "__main__":
     chat_memory = ChatbotMemory()
-    user_input = "Bonjour, comment allez-vous?"
-    bot_response = "Je vais bien, merci ! Et vous ?"
-    chat_memory.update_memory(user_input, bot_response)
+    
+    chat_memory.update_memory("Bonjour, comment allez-vous?", "Je vais bien, merci ! Et vous ?")
+    chat_memory.update_memory("Écris un poème de 300 mots", "Sous l'éclat argenté d'une lune éternelle... (poème)")
 
-    # Obtenir l'historique des conversations
     historique = chat_memory.get_memory()
-    print(historique)
-    
-    user_input = "écris un poème de 300 mots"
-    bot_response = "\
-    L’écho des mondes\
-    Sous l’éclat argenté d’une lune éternelle,\
-    Les murmures du vent dansent, doux et rebelles.\
-    Ils transportent des rêves, des promesses lointaines,\
-    Effleurant l’horizon de montagnes anciennes.\
-    \
-    Dans la forêt profonde, les arbres se confient,\
-    Leurs racines tissent des récits infinis.\
-    Chaque feuille, chaque branche, porte un secret du temps,\
-    Un souffle de sagesse, une parole en suspens.\
-    \
-    Les rivières murmurent en suivant leur chemin,\
-    Leurs eaux claires reflètent le bleu du matin.\
-    Elles emportent les peines, les joies et les cris,\
-    Traversant les âges, ignorant les périls.\
-    \
-    Le soleil se lève, doré, majestueux,\
-    Réveillant la terre d’un sommeil lumineux.\
-    Il peint des ombres dansantes sur les collines,\
-    Offrant à chaque instant des lueurs divines.\
-    \
-    Dans ce vaste univers où tout semble figé,\
-    Les étoiles veillent, telles des âmes égarées.\
-    Elles brillent de loin, gardiennes silencieuses,\
-    De secrets millénaires et de vies mystérieuses.\
-    \
-    Chaque battement de cœur, chaque souffle d’air,\
-    Est un fragment du tout, un lien universel.\
-    Nous marchons sur ce fil, entre ombre et lumière,\
-    Cherchant notre place dans cette grande sphère.\
-    \
-    Et dans le silence, là où tout se résout,\
-    L’écho des mondes résonne, doux mais flou.\
-    Il nous rappelle que l’immensité est en nous,\
-    Et que, dans chaque instant, réside l’infini goût.\
-    \
-    Le temps s’efface, les frontières se dissolvent,\
-    Dans ce ballet cosmique où les âmes s’envolent.\
-    Nous sommes poussière d’étoile, brève, éphémère,\
-    Mais dans l’éternité, notre essence prospère."
-    
-    chat_memory.update_memory(user_input, bot_response)
+    for entry in historique:
+        print(entry)