Update app.py

app.py

@@ -2,12 +2,36 @@ import gradio as gr
 import random
 import berkeleydb
 import pickle
+from sentence_transformers import SentenceTransformer, CrossEncoder, util
+
 
 class MariannaBot:
     def __init__(self):
         self.database = berkeleydb.hashopen("wiki_napoli_main.db", flag="c")
         self.database_legends = berkeleydb.hashopen("wiki_naples_leggende.db", flag="c")
+        self.db_keys = [key.decode("utf-8") for key, value in self.database.items()]
         self.reset_state()
+
+    def initialize_encoder(self):
+        """
+        Inizializza il modello encoder e pre-calcola gli embedding delle chiavi del database.
+        Questo metodo dovrebbe essere chiamato una sola volta all'avvio del bot.
+        """
+        try:
+            # Initialize the encoder model
+            encoder_model = "nickprock/sentence-bert-base-italian-uncased"
+            cross_encoder_model = "nickprock/cross-encoder-italian-bert-stsb"
+            self.encoder = SentenceTransformer(encoder_model)
+            self.cross_encoder = CrossEncoder(cross_encoder_model)
+            
+            # Pre-encode all database keys
+            self.db_keys_embeddings = self.encoder.encode(self.db_keys, convert_to_tensor=True)
+            
+            print(f"Encoder initialized with {len(self.db_keys)} keys.")
+            return True
+        except Exception as e:
+            print(f"Error initializing encoder: {str(e)}")
+        return False
     
     def reset_state(self):
         self.state = "initial"
@@ -20,9 +44,9 @@ class MariannaBot:
     def get_welcome_message(self):
         return """Ciao, benvenuto!\n\nSono Marianna, la testa di Napoli, in napoletano 'a capa 'e Napule, una statua ritrovata per caso nel 1594. \nAll'epoca del mio ritrovamento, si pensò che fossi una rappresentazione della sirena Partenope, dalle cui spoglie, leggenda narra, nacque la città di Napoli. In seguito, diversi studiosi riconobbero in me una statua della dea Venere, probabilmente collocata in uno dei tanti templi che si trovavano nella città in epoca tardo-romana, quando ancora si chiamava Neapolis.
         \nPosso raccontarti molte storie sulla città di Napoli e mostrarti le sue bellezze. \nC'è qualcosa in particolare che ti interessa?
-        \n(Rispondi con 'sì', 'no' o 'non so, scegli tu')"""
+        \n(Rispondi con 'sì', 'no' o 'non so')"""
     
-    def get_safe_example_keys(self, num_examples=5):
+    def get_safe_example_keys(self, num_examples=3):
         """Safely get example keys from the database."""
         try:
             keys = list(self.database.keys())
@@ -62,10 +86,34 @@ class MariannaBot:
     def handle_query(self, message):
         """Handle user queries by searching the database"""
         try:
-            for key, value in self.database.items():
-                decoded_key = key.decode("utf-8").lower()
-                if message == decoded_key:
-                    self.main_k.append(key.decode("utf-8"))
+            # Encode the user query
+            query_embedding = self.encoder.encode(message, convert_to_tensor=True)
+            
+            # Perform semantic search on the keys
+            semantic_hits = util.semantic_search(query_embedding, self.db_keys_embeddings, top_k=3)
+            semantic_hits = semantic_hits[0]
+
+            cross_inp = [(message, self.db_keys[hit['corpus_id']]) for hit in semantic_hits]
+            cross_scores = self.cross_encoder.predict(cross_inp)
+
+            reranked_hits = sorted(
+                [{'corpus_id': hit['corpus_id'], 'cross-score': score}
+                for hit, score in zip(semantic_hits, cross_scores)],
+                key=lambda x: x['cross-score'], reverse=True
+            )
+
+            best_hit = reranked_hits[0]
+            best_title = self.db_keys[best_hit['corpus_id']]
+            
+                
+            # Using your existing code for handling the best match
+            if best_title is not None:
+                best_title_bytes = best_title.encode("utf-8")  # Converti la stringa in bytes
+
+                if best_title_bytes in self.database:
+                    value = self.database[best_title_bytes] 
+                    key = best_title
+                    self.main_k.append(key)
                     self.state = "follow_up"
                     self.is_telling_stories = False
                     deserialized_value = pickle.loads(value)
@@ -73,8 +121,10 @@ class MariannaBot:
                     self.current_further_info_values = list(deserialized_value.get('further_info', {}).values())
                     self.current_index = 0
                     return f"{response}\n\nVuoi sapere altro su {self.main_k[-1]}? (sì/no)"
-            return "Mi dispiace, non ho informazioni riguardo a questa domanda. Prova a chiedermi qualcos'altro sulla città di Napoli."
-        except Exception:
+                else:
+                    return "Mi dispiace, non ho informazioni riguardo a questa domanda. Prova a chiedermi qualcos'altro sulla città di Napoli."
+        
+        except Exception as e:
             self.state = "initial"
             return "Mi dispiace, c'è stato un errore. Puoi riprovare con un'altra domanda?"
 
@@ -88,14 +138,14 @@ class MariannaBot:
             if message in ["sì", "si"]:
                 self.state = "query"
                 self.is_telling_stories = False
-                return "Di cosa vuoi sapere?"
+                return "Potresti dirmi di cosa vorresti sapere?"
             elif message == "no":
                 self.state = "end"
                 return "Va bene, grazie per aver parlato con me."
-            elif message == "non so, scegli te":
+            elif message == "non so":
                 return self.story_flow()
             else:
-                return "Scusa, non ho capito. Puoi rispondere con 'sì', 'no' o 'non so, scegli tu'."
+                return "Scusa, non ho capito. Puoi rispondere con 'sì', 'no' o 'non so'."
 
         elif self.state == "query":
             return self.handle_query(message)
@@ -128,6 +178,7 @@ class MariannaBot:
 
 def main():
     bot = MariannaBot()
+    bot.initialize_encoder()
 
     def update_chatbot(message, history):
         if not message.strip():
@@ -141,10 +192,10 @@ def main():
     
     with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
         with gr.Row():
-            gr.Markdown("## Parla con Marianna - 'La Testa di Napoli'")
+            gr.Markdown("## Chat con Marianna - 'La Testa di Napoli'")
             
         with gr.Row():
-            gr.Image("marianna-102.jpeg", 
+            gr.Image("/home/filippo/Scrivania/Marianna_head/Marianna_testa/marianna-102.jpeg", 
                     elem_id="marianna-image", 
                     width=250)