Update app.py

app.py

@@ -45,7 +45,7 @@ db = lancedb.connect("content 5/lancedb_data")
 table = db.open_table("sciencev4")
 
 def hybrid_search(text):
-    results = table.search(text, query_type="hybrid").limit(6).to_pandas()
+    results = table.search(text, query_type="hybrid").limit(5).to_pandas()
 
     document = []
     document_html = []
@@ -88,10 +88,21 @@ class pleiasBot:
                 eos_token_id=tokenizer.eos_token_id
             )
 
-            # Decode only the new tokens
+            # Decode the generated text
             generated_text = tokenizer.decode(output[0][len(input_ids[0]):])
             
-            generated_text = '<h2 style="text-align:center">Réponse</h3>\n<div class="generation">' + format_references(generated_text) + "</div>"
+            # Split the text into analysis and answer sections
+            parts = generated_text.split("<|source_analysis_end|>")
+            if len(parts) == 2:
+                analysis = parts[0].strip()
+                answer = parts[1].replace("<|answer_start|>", "").replace("<|answer_end|>", "").strip()
+                
+                # Format with section titles
+                formatted_text = f'<div class="section-title">Sources Analysis</div>\n\n{analysis}\n\n<div class="section-title">Answer</div>\n\n{answer}'
+            else:
+                formatted_text = generated_text
+            
+            generated_text = '<h2 style="text-align:center">Réponse</h3>\n<div class="generation">' + format_references(formatted_text) + "</div>"
             fiches_html = '<h2 style="text-align:center">Sources</h3>\n' + fiches_html
             return generated_text, fiches_html
 
@@ -102,7 +113,6 @@ class pleiasBot:
             return None, None
 
 def format_references(text):
-    # New ref format pattern
     ref_pattern = r'<ref name="([^"]+)">"([^"]+)"</ref>'
     
     parts = []
@@ -118,8 +128,8 @@ def format_references(text):
         ref_id = match.group(1)  # The source ID
         ref_text = match.group(2).strip()  # The reference text
         
-        # Create tooltip HTML
-        tooltip_html = f'<span class="tooltip"><strong>[{ref_number}]</strong><span class="tooltiptext"><strong>{ref_number}:</strong> {ref_text}</span></span>'
+        # Create tooltip HTML with source ID in bold
+        tooltip_html = f'<span class="tooltip">[{ref_number}]<span class="tooltiptext"><strong>{ref_id}</strong>: {ref_text}</span></span>'
         parts.append(tooltip_html)
         
         current_pos = match.end()
@@ -131,7 +141,7 @@ def format_references(text):
     return ''.join(parts)
 
 # Initialize the pleiasBot
-cassandre_bot = pleiasBot()
+pleias_bot = pleiasBot()
 
 # CSS for styling 
 css = """
@@ -195,10 +205,9 @@ css = """
     border-color: #fff transparent transparent transparent;
 }
 """
-
 # Gradio interface
 def gradio_interface(user_message):
-    response, sources = cassandre_bot.predict(user_message)
+    response, sources = pleias_bot.predict(user_message)
     return response, sources
 
 # Create Gradio app