Update app.py

app.py

@@ -18,7 +18,6 @@ def log_debug(message):
         full_message = f"[{timestamp}] {message}"
         debug_messages.append(full_message)
         print(full_message)  # Print to console
-        # Keep only the last 20 messages
         if len(debug_messages) > 20:
             debug_messages.pop(0)
         return "\n".join(debug_messages)
@@ -38,32 +37,40 @@ if not os.path.exists(MODEL_CACHE_DIR):
 log_debug("Initializing RAG model...")
 try:
     rag = RAGWithCitations(model_path_or_name=MODEL_CACHE_DIR)
-    log_debug("✅ Model loaded successfully")
+    
+    # Fix the warnings by properly configuring generation parameters
     if hasattr(rag, "model"):
-        log_debug(f"Model device: {rag.model.device}")
-        log_debug(f"Model config: {rag.model.config}")
+        # Configure tokenizer
+        if hasattr(rag, "tokenizer"):
+            if rag.tokenizer.pad_token is None:
+                rag.tokenizer.pad_token = rag.tokenizer.eos_token
+            rag.tokenizer.padding_side = "left"  # For batch generation
+            
+        # Configure model generation settings
+        rag.model.config.pad_token_id = rag.tokenizer.pad_token_id
+        rag.model.generation_config.pad_token_id = rag.tokenizer.pad_token_id
+        
+        # Fix the do_sample/top_p warning
+        rag.model.generation_config.do_sample = True
+        rag.model.generation_config.top_p = 0.95  # Explicitly set to match warning
+        
+        # Configure attention mask handling
+        rag.model.config.use_cache = True
+        
+        log_debug("✅ Model loaded successfully with configuration:")
+        log_debug(f" - Pad token: {rag.tokenizer.pad_token} (ID: {rag.tokenizer.pad_token_id})")
+        log_debug(f" - Generation config: {rag.model.generation_config}")
+        
 except Exception as e:
     log_debug(f"❌ Model initialization failed: {str(e)}")
     raise
 
-# Model configuration
-if hasattr(rag, "tokenizer"):
-    rag.tokenizer.pad_token = rag.tokenizer.eos_token
-if hasattr(rag, "model"):
-    rag.model.config.pad_token_id = rag.tokenizer.eos_token_id
-    rag.model.generation_config.do_sample = True
-    rag.model.config.use_cache = True
-
 def extract_text_from_pdf_url(url, debug_state):
     """Extract text from PDF with debug logging"""
     debug_state = log_debug(f"📄 Fetching PDF: {url[:60]}...")
     try:
-        start_time = time.time()
         response = requests.get(url, timeout=30)
         response.raise_for_status()
-        load_time = time.time() - start_time
-        debug_state = log_debug(f"⏳ PDF downloaded in {load_time:.2f}s (size: {len(response.content)/1024:.1f}KB)")
-        
         doc = fitz.open(stream=response.content, filetype="pdf")
         text = ""
         for page in doc:
@@ -88,12 +95,14 @@ def generate_answer(query, pdf_urls_str, debug_state=""):
         sources = []
         feedback = "### PDF Load Report:\n"
         
-        debug_state = log_debug(f"Processing {len(pdf_urls)} PDF URLs...")
-        
         for url in pdf_urls:
             text, debug_state = extract_text_from_pdf_url(url, debug_state)
             if not text.startswith("[Error"):
-                sources.append({"text": text, "metadata": {"source": url}})
+                sources.append({
+                    "text": text,
+                    "metadata": {"source": url},
+                    "attention_mask": [1] * len(text.split())  # Simple attention mask
+                })
                 feedback += f"- ✅ Loaded: {url[:80]}\n"
             else:
                 feedback += f"- ❌ Failed: {url[:80]}\n"
@@ -103,20 +112,20 @@ def generate_answer(query, pdf_urls_str, debug_state=""):
             return feedback + "\nNo valid PDFs processed", debug_state
         
         debug_state = log_debug(f"🧠 Generating answer using {len(sources)} sources...")
-        start_time = time.time()
         
         try:
-            response = rag.generate(query, sources)
-            gen_time = time.time() - start_time
-            debug_state = log_debug(f"⚡ Generation completed in {gen_time:.2f}s")
+            # Generate with proper attention handling
+            response = rag.generate(
+                query, 
+                sources,
+                attention_mask=True  # Ensure attention masks are used
+            )
             
             answer = response.get('raw_response', 'No response generated')
             backend = response.get('backend_used', 'unknown')
             
-            debug_state = log_debug(f"💡 Answer preview: {answer[:200]}...")
-            debug_state = log_debug(f"🛠️ Backend used: {backend}")
-            
-            full_output = f"{feedback}\n\n### Answer:\n{answer}\n\n_Generated in {gen_time:.2f}s using {backend}_"
+            debug_state = log_debug(f"💡 Answer generated using {backend}")
+            full_output = f"{feedback}\n\n### Answer:\n{answer}\n\n_Generated using {backend}_"
             return full_output, debug_state
             
         except Exception as e:
@@ -141,13 +150,9 @@ with gr.Blocks(title="Pleias RAG QA", css="""
         padding: 10px;
         border-radius: 5px;
     }
-    .debug-title {
-        font-weight: bold;
-        margin-bottom: 5px;
-    }
 """) as demo:
     
-    gr.Markdown("## Retrieval Generation from PDF files with a 350MB Pocket Size Model from Pleias")
+    gr.Markdown("# Retrieval Generation from PDF files with a 350MB Pocket Size Model from Pleias")
     
     with gr.Row():
         with gr.Column():
@@ -159,15 +164,13 @@ with gr.Blocks(title="Pleias RAG QA", css="""
         with gr.Column():
             output = gr.Markdown(label="Model Response")
             if DEBUG:
-                gr.Markdown("### Debug Console", elem_classes=["debug-title"])
                 debug_console = gr.Textbox(
-                    label="",
+                    label="Debug Console",
                     interactive=False,
                     lines=15,
                     elem_classes=["debug-console"]
                 )
     
-    # Handle submission
     submit_btn.click(
         fn=generate_answer,
         inputs=[question, pdf_urls] + ([debug_console] if DEBUG else []),