Initial Commit

app.py

@@ -0,0 +1,578 @@
+import os
+import time
+from typing import List, Dict, Any, Optional
+from concurrent.futures import ThreadPoolExecutor
+
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from Bio import Entrez
+import traceback
+import pandas as pd
+
+
+# Cache last-created agent to avoid reloading the model on every call
+_CACHED_AGENT_KEY = None
+_CACHED_AGENT = None
+
+# Also cache model/tokenizer per device to prevent repeated downloads
+_MODEL_CACHE: Dict[str, Dict[str, Any]] = {}
+
+
+MODEL_NAME = "hkust-nlp/WebExplorer-8B"
+
+
+def _get_hf_components(device_str: str) -> Dict[str, Any]:
+    """Load and cache tokenizer/model for the requested device string."""
+    if device_str in _MODEL_CACHE:
+        return _MODEL_CACHE[device_str]
+    
+    print(f"Loading model for device: {device_str}")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
+    
+    # Configure 4-bit quantization for much faster loading and inference (with safe fallback)
+    if torch.cuda.is_available():
+        try:
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4"
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                quantization_config=quantization_config,
+                device_map="auto",
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+            )
+        except Exception as e:
+            print(f"4-bit load failed, falling back to standard half precision: {e}")
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+            )
+    else:
+        # CPU fallback (slower)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            device_map="auto",
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True,
+        )
+    
+    # Set padding token if not set
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        model.config.pad_token_id = tokenizer.eos_token_id
+    
+    print(f"Model loaded successfully on {device_str}")
+    _MODEL_CACHE[device_str] = {"tokenizer": tokenizer, "model": model}
+    return _MODEL_CACHE[device_str]
+
+
+class LocalWebExplorerAgent:
+    """Optimized medical research agent with PubMed integration."""
+
+    def __init__(self, search_targets: List[str], use_cpu: bool):
+        self.search_targets = search_targets
+        self.device_str = "cpu" if use_cpu else ("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Configure Entrez from environment variables if present
+        Entrez.email = os.getenv("ENTREZ_EMAIL", "harshini.kalvakuntla@gmail.com")
+        Entrez.api_key = os.getenv("ENTREZ_API_KEY","e87e8f21aeaa01cdd5690c52e8a4f5336008")
+
+        comps = _get_hf_components(self.device_str)
+        self.tokenizer = comps["tokenizer"]
+        self.model = comps["model"]
+        
+        # Cache for search results to avoid redundant API calls
+        self.search_cache: Dict[str, List[Dict[str, str]]] = {}
+
+    def _needs_search(self, query: str) -> bool:
+        """Determine if external search is needed."""
+        lowered = query.lower()
+        trigger_terms = [
+            "treatment", "survival", "trial", "latest", "guideline", 
+            "therapy", "diagnosis", "prognosis", "rate", "statistic",
+            "study", "research", "clinical", "evidence"
+        ]
+        return any(term in lowered for term in trigger_terms)
+
+    def _extract_diagnosis(self, query: str) -> str:
+        """Extract medical condition from query."""
+        query_lower = query.lower()
+        
+        # Common conditions mapping
+        conditions = {
+            "lung": "lung cancer",
+            "pancreatic": "pancreatic cancer",
+            "breast": "breast cancer",
+            "colon": "colorectal cancer",
+            "prostate": "prostate cancer",
+            "melanoma": "melanoma",
+            "diabetes": "diabetes mellitus",
+            "heart failure": "heart failure",
+            "hypertension": "hypertension",
+        }
+        
+        for key, value in conditions.items():
+            if key in query_lower:
+                return value
+        
+        return "general medical condition"
+
+    def _pubmed_search(self, diagnosis: str) -> List[Dict[str, str]]:
+        """Search PubMed with caching."""
+        # Check cache first
+        if diagnosis in self.search_cache:
+            return self.search_cache[diagnosis]
+        
+        if not Entrez.email or Entrez.email == "user@example.com":
+            # Return empty if no valid email configured
+            return []
+        
+        try:
+            query = f"{diagnosis} treatment guidelines[Title/Abstract] OR {diagnosis} clinical practice[Title/Abstract]"
+            handle = Entrez.esearch(db="pubmed", term=query, retmax=3, sort="relevance")
+            record = Entrez.read(handle)
+            handle.close()
+            
+            ids = record.get("IdList", [])
+            results: List[Dict[str, str]] = []
+            
+            if ids:
+                # Fetch summaries in batch
+                fetch = Entrez.esummary(db="pubmed", id=",".join(ids), retmode="xml")
+                summary_list = Entrez.read(fetch)
+                fetch.close()
+                
+                for summary in summary_list:
+                    pmid = summary.get("Id", "")
+                    title = summary.get("Title", "No title")
+                    results.append({
+                        "pmid": str(pmid),
+                        "title": title,
+                        "url": f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/"
+                    })
+            
+            # Cache results
+            self.search_cache[diagnosis] = results
+            return results
+            
+        except Exception as e:
+            print(f"PubMed search error: {e}")
+            return []
+
+    def _fetch_abstracts(self, pmids: List[str]) -> str:
+        """Fetch abstracts in parallel for speed."""
+        if not Entrez.email or not pmids:
+            return ""
+        
+        def fetch_single(pmid: str) -> str:
+            try:
+                fetch = Entrez.efetch(db="pubmed", id=pmid, rettype="abstract", retmode="text")
+                content = fetch.read()
+                fetch.close()
+                
+                if isinstance(content, bytes):
+                    content = content.decode('utf-8', errors='ignore')
+                return content
+            except Exception as e:
+                print(f"Error fetching abstract for PMID {pmid}: {e}")
+                return ""
+        
+        # Use ThreadPoolExecutor for parallel fetching
+        with ThreadPoolExecutor(max_workers=3) as executor:
+            abstracts = list(executor.map(fetch_single, pmids))
+        
+        return "\n\n".join([a for a in abstracts if a])
+
+    def _generate(self, prompt: str, max_new_tokens: int = 200) -> str:
+        """Optimized generation with proper settings."""
+        inputs = self.tokenizer(
+            prompt, 
+            return_tensors="pt",
+            truncation=True,
+            max_length=1024  # Limit input length for speed
+        ).to(self.model.device)
+        
+        with torch.inference_mode():  # Faster than torch.no_grad()
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,  # Greedy decoding is fastest
+                num_beams=1,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                use_cache=True,  # KV cache for speed
+            )
+        
+        # Decode only the generated tokens
+        generated_ids = outputs[0][inputs.input_ids.shape[1]:]
+        return self.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+
+    def execute_query(self, query: str, max_turns: int = 3) -> Dict[str, Any]:
+        """Execute a single query with optimized flow."""
+        turns: List[Dict[str, Any]] = []
+        timestamp = int(time.time())
+
+        # Extract diagnosis
+        diagnosis = self._extract_diagnosis(query)
+
+        # Turn 1: decision
+        needs_search = self._needs_search(query)
+        turns.append({
+            "turn": 1,
+            "action_decision": "search" if needs_search else "reason",
+            "tool_calls": [],
+        })
+
+        retrieved_docs: List[Dict[str, str]] = []
+        abstracts = ""
+        
+        # Turn 2: search if needed
+        if needs_search and len(turns) < max_turns:
+            retrieved_docs = self._pubmed_search(diagnosis)
+            turns.append({
+                "turn": 2,
+                "action_decision": "search",
+                "tool_calls": [{
+                    "tool": "pubmed.search",
+                    "args": {"diagnosis": diagnosis},
+                    "results": [f"PMID {d['pmid']}: {d['title']}" for d in retrieved_docs],
+                }],
+            })
+            
+            # Fetch abstracts if we have PMIDs
+            if retrieved_docs:
+                pmids = [d["pmid"] for d in retrieved_docs]
+                abstracts = self._fetch_abstracts(pmids)
+
+        # Turn 3: Generate answer
+        prompt = self._build_prompt(query, diagnosis, abstracts)
+        answer_text = self._generate(prompt, max_new_tokens=200)
+        
+        turns.append({
+            "turn": len(turns) + 1,
+            "action_decision": "reason",
+            "tool_calls": [],
+            "response": answer_text[:100] + "..."
+        })
+
+        # Add disclaimer and sources
+        answer_text = self._format_answer(answer_text, query, retrieved_docs)
+
+        return {
+            "model_loaded": True,
+            "final_answer": answer_text,
+            "turns": turns,
+            "total_turns": len(turns),
+            "timestamp": timestamp,
+        }
+
+    def _build_prompt(self, query: str, diagnosis: str, abstracts: str) -> str:
+        """Build optimized prompt."""
+        if abstracts:
+            return (
+                f"Answer this medical question based on the research below.\n\n"
+                f"Question: {query}\n\n"
+                f"Research on {diagnosis}:\n{abstracts[:1500]}\n\n"  # Limit context
+                f"Provide a clear, concise summary of current treatments and outcomes."
+            )
+        else:
+            return (
+                f"Answer this medical question concisely and accurately.\n\n"
+                f"Question: {query}\n\n"
+                f"Provide evidence-based information in plain language."
+            )
+
+    def _format_answer(self, answer: str, query: str, docs: List[Dict[str, str]]) -> str:
+        """Format answer with disclaimer and sources."""
+        # Add medical disclaimer
+        medical_terms = ["cancer", "disease", "diabetes", "treatment", "diagnosis", "therapy"]
+        if any(term in query.lower() for term in medical_terms):
+            answer += "\n\n**Disclaimer:** This is educational information only. Always consult a healthcare professional for medical advice."
+        
+        # Add sources
+        if docs:
+            answer += "\n\n**Sources:**\n" + "\n".join(
+                f"- [{d['title']}]({d['url']})" for d in docs
+            )
+        
+        return answer
+
+    def execute_batch(self, queries: List[str], max_turns: int = 3, progress_callback=None) -> List[Dict[str, Any]]:
+        """Process multiple queries with progress tracking."""
+        results = []
+        total = len(queries)
+        
+        for idx, query in enumerate(queries):
+            if progress_callback:
+                progress_callback((idx + 1) / total, desc=f"Processing query {idx + 1}/{total}")
+            
+            try:
+                result = self.execute_query(query, max_turns=max_turns)
+                results.append(result)
+            except Exception as e:
+                print(f"Error processing query '{query}': {e}")
+                results.append({
+                    "model_loaded": False,
+                    "final_answer": f"Error: {str(e)}",
+                    "turns": [],
+                    "total_turns": 0,
+                    "timestamp": int(time.time()),
+                    "error": str(e)
+                })
+        
+        return results
+
+
+DEFAULT_TARGETS = [
+    'nih.gov', 'cdc.gov', 'fda.gov', 'clinicaltrials.gov', 'medlineplus.gov',
+    'who.int', 'cancerresearchuk.org', 'esmo.org', 'cancer.org', 'cancer.net',
+    'mayoclinic.org', 'mdanderson.org', 'mskcc.org', 'dana-farber.org',
+    'uptodate.com', 'ncbi.nlm.nih.gov', 'healthline.com',
+]
+
+
+def get_agent(search_targets: List[str], use_cpu: bool) -> LocalWebExplorerAgent:
+    """Get or create cached agent."""
+    global _CACHED_AGENT_KEY, _CACHED_AGENT
+    key = (tuple(sorted(search_targets)), use_cpu)
+    if _CACHED_AGENT is not None and _CACHED_AGENT_KEY == key:
+        return _CACHED_AGENT
+    _CACHED_AGENT = LocalWebExplorerAgent(search_targets=search_targets, use_cpu=use_cpu)
+    _CACHED_AGENT_KEY = key
+    return _CACHED_AGENT
+
+
+def run_query(query: str, domain_scope: str, device_choice: str, max_turns: int, fast_mode: bool, progress=gr.Progress()):
+    """Run a single query with progress tracking."""
+    if not query or not query.strip():
+        return "Please enter a query.", {}
+
+    progress(0, desc="Loading model...")
+    use_cpu = device_choice == "CPU"
+    targets = DEFAULT_TARGETS if domain_scope == "Medical (Trusted sources only)" else []
+    
+    try:
+        agent = get_agent(targets, use_cpu=use_cpu)
+        progress(0.2, desc="Processing query...")
+        
+        if fast_mode:
+            # Fast path: skip PubMed and generate a concise answer with fewer tokens
+            agent._needs_search = lambda q: False  # bypass search
+            result = agent.execute_query(query.strip(), max_turns=1)
+            # Truncate final answer if too long
+            if result.get('final_answer'):
+                result['final_answer'] = result['final_answer'][:1200]
+        else:
+            result = agent.execute_query(query.strip(), max_turns=max_turns)
+        
+        progress(1.0, desc="Complete!")
+        final_answer = result.get('final_answer', '')
+        
+        mini_trace = {
+            'model_loaded': result.get('model_loaded'),
+            'turns': result.get('turns', []),
+            'total_turns': result.get('total_turns'),
+            'timestamp': result.get('timestamp'),
+            'fast_mode': fast_mode,
+        }
+        return final_answer, mini_trace
+        
+    except Exception as e:
+        tb = traceback.format_exc()
+        print("\n===== ERROR IN run_query =====\n", tb, "\n==============================\n")
+        return f"Error: {str(e)}", {"error": str(e), "traceback": tb}
+
+
+def process_batch_file(file, domain_scope: str, device_choice: str, max_turns: int, progress=gr.Progress()):
+    """Process batch file with queries."""
+    if file is None:
+        return "Please upload a file.", None
+    
+    progress(0, desc="Reading file...")
+    
+    try:
+        # Read queries
+        if file.name.endswith('.csv'):
+            df = pd.read_csv(file.name)
+            if 'query' in df.columns:
+                queries = df['query'].tolist()
+            elif 'question' in df.columns:
+                queries = df['question'].tolist()
+            else:
+                queries = df.iloc[:, 0].tolist()
+        elif file.name.endswith('.txt'):
+            with open(file.name, 'r', encoding='utf-8') as f:
+                queries = [line.strip() for line in f if line.strip()]
+        else:
+            return "Please upload a CSV or TXT file.", None
+        
+        if not queries:
+            return "No queries found in file.", None
+        
+        progress(0.1, desc=f"Found {len(queries)} queries. Loading model...")
+        
+        use_cpu = device_choice == "CPU"
+        targets = DEFAULT_TARGETS if domain_scope == "Medical (Trusted sources only)" else []
+        agent = get_agent(targets, use_cpu=use_cpu)
+        
+        # Process batch
+        results = agent.execute_batch(
+            queries, 
+            max_turns=max_turns,
+            progress_callback=lambda p, desc: progress(0.1 + p * 0.9, desc=desc)
+        )
+        
+        # Create results dataframe
+        results_data = []
+        for query, result in zip(queries, results):
+            results_data.append({
+                'Query': query,
+                'Answer': result.get('final_answer', 'Error'),
+                'Total Turns': result.get('total_turns', 0),
+                'Success': result.get('model_loaded', False),
+            })
+        
+        results_df = pd.DataFrame(results_data)
+        
+        # Save results
+        output_path = f"batch_results_{int(time.time())}.csv"
+        results_df.to_csv(output_path, index=False)
+        
+        progress(1.0, desc="Complete!")
+        
+        success_count = sum(r.get('model_loaded', False) for r in results)
+        summary = (
+            f"✅ Processed {len(queries)} queries\n\n"
+            f"📊 Success rate: {success_count}/{len(results)}\n\n"
+            f"💾 Results saved to: `{output_path}`"
+        )
+        
+        return summary, results_df
+        
+    except Exception as e:
+        tb = traceback.format_exc()
+        print("\n===== ERROR IN process_batch_file =====\n", tb, "\n==============================\n")
+        return f"Error processing file: {e}", None
+
+
+# Gradio Interface
+with gr.Blocks(title="WebExplorer-8B Medical Research", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🔬 WebExplorer-8B Medical Research Assistant
+    Ask medical questions or process multiple queries in batch. Powered by AI and PubMed research.
+    """)
+    
+    with gr.Tabs():
+        with gr.Tab("💬 Single Query"):
+            with gr.Row():
+                query = gr.Textbox(
+                    label="Medical Question", 
+                    lines=3, 
+                    placeholder="e.g., What are the treatment options for Type 2 diabetes?",
+                    scale=4
+                )
+            
+            with gr.Row():
+                domain_scope = gr.Radio(
+                    choices=["Medical (Trusted sources only)", "All sources"], 
+                    value="Medical (Trusted sources only)", 
+                    label="Source Scope",
+                    scale=2
+                )
+                device = gr.Radio(
+                    choices=["GPU", "CPU"], 
+                    value="GPU", 
+                    label="Device",
+                    scale=1
+                )
+                max_turns = gr.Slider(
+                    minimum=1, maximum=5, value=2, step=1, 
+                    label="Max Research Depth",
+                    scale=1
+                )
+                fast_mode = gr.Checkbox(value=True, label="Fast mode (skip PubMed, shorter answer)")
+
+            submit = gr.Button("🔍 Research", variant="primary", size="lg")
+
+            answer = gr.Markdown(label="Answer", height=300)
+            trace = gr.Json(label="Execution Trace", visible=False)
+
+            gr.Markdown("### 📚 Example Questions")
+            gr.Examples(
+                examples=[
+                    ["What are the survival rates for stage IV pancreatic cancer?"],
+                    ["How is Type 2 diabetes diagnosed and treated?"],
+                    ["What are the latest immunotherapy options for melanoma?"],
+                    ["What are the risk factors for colorectal cancer?"],
+                ],
+                inputs=[query],
+            )
+
+            submit.click(
+                run_query, 
+                inputs=[query, domain_scope, device, max_turns, fast_mode], 
+                outputs=[answer, trace]
+            )
+        
+        with gr.Tab("📊 Batch Processing"):
+            gr.Markdown("""
+            ### Process Multiple Queries
+            Upload a **CSV** (with 'query' column) or **TXT** file (one query per line).
+            """)
+            
+            batch_file = gr.File(
+                label="Upload File",
+                file_types=['.csv', '.txt'],
+                scale=2
+            )
+            
+            with gr.Row():
+                batch_domain = gr.Radio(
+                    choices=["Medical (Trusted sources only)", "All sources"], 
+                    value="Medical (Trusted sources only)", 
+                    label="Source Scope"
+                )
+                batch_device = gr.Radio(
+                    choices=["GPU", "CPU"], 
+                    value="GPU", 
+                    label="Device"
+                )
+                batch_turns = gr.Slider(
+                    minimum=1, maximum=5, value=2, step=1, 
+                    label="Max Research Depth"
+                )
+            
+            batch_submit = gr.Button("🚀 Process Batch", variant="primary", size="lg")
+            
+            batch_status = gr.Markdown(label="Status")
+            batch_results = gr.Dataframe(label="Results Preview", max_height=400)
+            
+            batch_submit.click(
+                process_batch_file,
+                inputs=[batch_file, batch_domain, batch_device, batch_turns],
+                outputs=[batch_status, batch_results]
+            )
+
+    gr.Markdown("""
+    ---
+    **Note:** Configure `ENTREZ_EMAIL` environment variable for PubMed access. 
+    GPU recommended for faster processing (2-5s vs 30-60s on CPU).
+    """)
+
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", "7860"))
+    demo.launch(
+        server_name="0.0.0.0", 
+        server_port=port, 
+        show_api=False,
+        share=False
+    )

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+torch
+transformers
+accelerate
+bitsandbytes
+biopython
+pandas