Update app.py

app.py

@@ -45,15 +45,17 @@ log_debug("Initializing RAG model...")
 try:
     #rag = RAGWithCitations(model_path_or_name=MODEL_CACHE_DIR)
     rag = RAGWithCitations(
-        model_path_or_name="1b_rag",
-        max_tokens=2048,        # Maximum tokens to generate (default: 2048)
-        temperature=0.0,        # Sampling temperature (default: 0.0)
-        top_p=0.95,             # Nucleus sampling parameter (default: 0.95)
-        repetition_penalty=1.0, # Penalty to reduce repetition (default: 1.0)
-        trust_remote_code=True, # Whether to trust remote code (default: True)
-        hf_token=os.getenv("HF_TOKEN")#,  # Required for downloading predefined models
-       # models_dir=MODEL_CACHE_DIR   # Custom directory for downloaded models
+        model_path_or_name="PleIAs/Pleias-RAG-350M"
     )
+#        model_path_or_name="1b_rag",
+#        max_tokens=2048,        # Maximum tokens to generate (default: 2048)
+#        temperature=0.0,        # Sampling temperature (default: 0.0)
+#        top_p=0.95,             # Nucleus sampling parameter (default: 0.95)
+#        repetition_penalty=1.0, # Penalty to reduce repetition (default: 1.0)
+#        trust_remote_code=True, # Whether to trust remote code (default: True)
+#        hf_token=os.getenv("HF_TOKEN")#,  # Required for downloading predefined models
+       # models_dir=MODEL_CACHE_DIR   # Custom directory for downloaded models
+#    )
     
     # Fix the warnings by properly configuring generation parameters
 #    if hasattr(rag, "model"):
@@ -86,39 +88,39 @@ except Exception as e:
 ## Let's a do simple test from the doc --
 
 # Define query and sources
-query = "What is the capital of France?"
-log_debug(f"🔍 Test Query: {query}")
-
-sources = [
-    {
-        "text": "Paris is the capital and most populous city of France.",
-        "metadata": {"source": "Geographic Encyclopedia", "reliability": "high"}
-    },
-    {
-        "text": "The Eiffel Tower is located in Paris, France.",
-        "metadata": {"source": "Travel Guide", "year": 2020}
-    }
-]
-log_debug("📄 Test Sources loaded successfully.")
+#query = "What is the capital of France?"
+#log_debug(f"🔍 Test Query: {query}")
+
+#sources = [
+#    {
+#        "text": "Paris is the capital and most populous city of France.",
+#        "metadata": {"source": "Geographic Encyclopedia", "reliability": "high"}
+#    },
+#    {
+#        "text": "The Eiffel Tower is located in Paris, France.",
+#        "metadata": {"source": "Travel Guide", "year": 2020}
+#    }
+#]
+#log_debug("📄 Test Sources loaded successfully.")
 
 # Generate a response
-try:
-    log_debug("🧠 Test rag model on simple example...")
-    rag1 = RAGWithCitations(
-      model_path_or_name="PleIAs/Pleias-RAG-350M"
-    )
-    response = rag1.generate(query, 
-                            sources #,
+#try:
+#    log_debug("🧠 Test rag model on simple example...")
+#    rag1 = RAGWithCitations(
+#      model_path_or_name="PleIAs/Pleias-RAG-350M"
+#    )
+#    response = rag1.generate(query, 
+#                            sources #,
                            # do_sample=True,  # Enable sampling
                            # top_p=0.95,     # Set top_p for nucleus sampling
                            # pad_token_id=rag.tokenizer.eos_token_id,  # Set pad_token_id to eos_token_id
                            # attention_mask=None  # Ensure attention_mask is passed if needed
-                           )
-    log_debug("✅ Test Answer generated successfully.")
-    log_debug(response["processed"]["clean_answer"])
-except Exception as e:
-    log_debug(f"❌ Test Answer generation failed: {str(e)}")
-    raise
+#                           )
+#    log_debug("✅ Test Answer generated successfully.")
+#    log_debug(response["processed"]["clean_answer"])
+#except Exception as e:
+#    log_debug(f"❌ Test Answer generation failed: {str(e)}")
+#    raise
 
 
 
@@ -180,14 +182,12 @@ def generate_answer(query, pdf_urls_str, debug_state=""):
             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')
+            answer = response["processed"]["clean_answer"]
             
             debug_state = log_debug(f"💡 Answer preview: {answer[:200]}...")
-            debug_state = log_debug(f"🛠️ Backend used: {backend}")
+            debug_state = log_debug(f"🛠️ Generated in {gen_time:.2f}s")
             
-            full_output = f"{feedback}\n\n### Answer:\n{answer}\n\n_Generated in {gen_time:.2f}s using {backend}_"
-            return full_output, debug_state
+            return answer, debug_state
             
         except Exception as e:
             error_msg = f"❌ Generation error: {str(e)}"