Update app.py

app.py

@@ -6,8 +6,6 @@ import time
 import traceback
 from huggingface_hub import snapshot_download
 from pleias_rag_interface import RAGWithCitations
-from fastapi import FastAPI
-from fastapi.staticfiles import StaticFiles
 
 # Debugging setup
 DEBUG = True
@@ -20,6 +18,7 @@ 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)
@@ -36,38 +35,50 @@ if not os.path.exists(MODEL_CACHE_DIR):
     log_debug("Downloading model...")
     snapshot_download(repo_id=MODEL_REPO, local_dir=MODEL_CACHE_DIR)
 
+
 log_debug("Initializing RAG model...")
 try:
     rag = RAGWithCitations(model_path_or_name=MODEL_CACHE_DIR)
-
-    # Configure tokenizer and model properly
+    
+    # Fix the warnings by properly configuring generation parameters
     if hasattr(rag, "model"):
+        # Configure tokenizer
         if hasattr(rag, "tokenizer"):
-            # Set pad token if not set
             if rag.tokenizer.pad_token is None:
                 rag.tokenizer.pad_token = rag.tokenizer.eos_token
-            rag.tokenizer.padding_side = "left"
-
-        # Configure generation settings
+            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
-
-        log_debug("✅ Model loaded with configuration:")
-        log_debug(f"Pad token: {rag.tokenizer.pad_token}")
-        log_debug(f"Generation config: {rag.model.generation_config}")
-
+        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
 
-def extract_text_from_pdf_url(url):
+
+
+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:
@@ -79,65 +90,64 @@ def extract_text_from_pdf_url(url):
         debug_state = log_debug(error_msg)
         return f"[Error loading PDF: {str(e)}]", debug_state
 
-def generate_answer(query, pdf_urls_str):
+def generate_answer(query, pdf_urls_str, debug_state=""):
     """Main processing function with debug output"""
     try:
         debug_state = log_debug(f"🔍 New query: {query}")
-
+        
         if not query or not pdf_urls_str:
             debug_state = log_debug("❌ Missing question or PDF URLs")
             return "Please provide both inputs", debug_state
-
+        
         pdf_urls = [url.strip() for url in pdf_urls_str.strip().split("\n") if url.strip()]
         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)
+            text, debug_state = extract_text_from_pdf_url(url, debug_state)
             if not text.startswith("[Error"):
-                # Create source document with metadata
-                sources.append({
-                    "text": text,
-                    "metadata": {"source": url}
-                })
+                sources.append({"text": text, "metadata": {"source": url}})
                 feedback += f"- ✅ Loaded: {url[:80]}\n"
             else:
                 feedback += f"- ❌ Failed: {url[:80]}\n"
-
+        
         if not sources:
             debug_state = log_debug("❌ No valid PDFs processed")
             return feedback + "\nNo valid PDFs processed", debug_state
-
+        
         debug_state = log_debug(f"🧠 Generating answer using {len(sources)} sources...")
-
+        start_time = time.time()
+        
         try:
-            # Generate without attention_mask parameter
             response = rag.generate(query, sources)
-
+            gen_time = time.time() - start_time
+            debug_state = log_debug(f"⚡ Generation completed in {gen_time:.2f}s")
+            
             answer = response.get('raw_response', 'No response generated')
             backend = response.get('backend_used', 'unknown')
-
-            debug_state = log_debug(f"💡 Answer generated using {backend}")
-            full_output = f"{feedback}\n\n### Answer:\n{answer}\n\n_Generated using {backend}_"
+            
+            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}_"
             return full_output, debug_state
-
+            
         except Exception as e:
             error_msg = f"❌ Generation error: {str(e)}"
             debug_state = log_debug(error_msg)
             debug_state = log_debug(traceback.format_exc())
             return feedback + f"\n\n❌ Error: {str(e)}", debug_state
-
+            
     except Exception as e:
         error_msg = f"❌ System error: {str(e)}"
         debug_state = log_debug(error_msg)
         debug_state = log_debug(traceback.format_exc())
         return error_msg, debug_state
 
-# Create the FastAPI app
-app = FastAPI()
-
-# Mount Gradio app
-gradio_app = gr.Blocks(title="Pleias RAG QA", css="""
+# Create the Gradio interface
+with gr.Blocks(title="Pleias RAG QA", css="""
     .debug-console {
         font-family: monospace;
         max-height: 400px;
@@ -146,38 +156,44 @@ gradio_app = gr.Blocks(title="Pleias RAG QA", css="""
         padding: 10px;
         border-radius: 5px;
     }
-""")
-
-with gradio_app:
+    .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")
-
+    
     with gr.Row():
         with gr.Column():
             question = gr.Textbox(label="Your Question", placeholder="What is this document about?")
-            pdf_urls = gr.Textbox(lines=5, label="PDF URLs (one per line)",
+            pdf_urls = gr.Textbox(lines=5, label="PDF URLs (one per line)", 
                                 placeholder="https://example.com/doc1.pdf")
             submit_btn = gr.Button("Submit", variant="primary")
-
+        
         with gr.Column():
             output = gr.Markdown(label="Model Response")
             if DEBUG:
+                gr.Markdown("### Debug Console", elem_classes=["debug-title"])
                 debug_console = gr.Textbox(
-                    label="Debug Console",
+                    label="",
                     interactive=False,
                     lines=15,
                     elem_classes=["debug-console"]
                 )
-
+    
+    # Handle submission
     submit_btn.click(
         fn=generate_answer,
-        inputs=[question, pdf_urls],
-        outputs=[output, debug_console],
+        inputs=[question, pdf_urls] + ([debug_console] if DEBUG else []),
+        outputs=[output, debug_console] if DEBUG else [output],
     )
 
-# Mount the Gradio app to FastAPI
-app.mount("/", gradio_app)
-
 if __name__ == "__main__":
     log_debug("🚀 Launching interface...")
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.launch(
+        server_port=7860,
+        server_name="0.0.0.0",
+        show_error=True,
+        debug=DEBUG
+    )