Update app.py

app.py

@@ -1,670 +1,260 @@
-import os
-import re
-from pathlib import Path
-import time
 import threading
-import torch
-import pandas as pd
+import time
+
 import gradio as gr
-from difflib import SequenceMatcher
-import unsloth
+import torch
 from unsloth import FastLanguageModel
-from peft import PeftModel
-from huggingface_hub import snapshot_download
-
-# --- 1. Configuration ---
-BASE_MODEL_ID = "unsloth/Meta-Llama-3.1-8B-Instruct"
-IT_ADAPTER_PATH = "./DA-IT"
-DPO_ADAPTER_PATH = "./DA-DPO"
-MAX_OPTIONS = 8
 
-DATA_PATH = os.path.join("Testing MCQA data", "Decarbonization_MCQA.csv")
 
-# --- 2. Load dataset ---
-try:
-    MCQA_DF = pd.read_csv(DATA_PATH, encoding="utf-8")
-except UnicodeDecodeError:
-    MCQA_DF = pd.read_csv(DATA_PATH, encoding="latin1")
+# ============================================================
+# Configuration
+# ============================================================
+BASE_MODEL_ID = "unsloth/Meta-Llama-3.1-8B-Instruct"
+LOAD_IN_4BIT = False
+
+DEFAULT_SYSTEM_PROMPT = """
+You are a scientific summarization assistant.
+
+Summarize only the supplied document.
+Do not continue the source text.
+Do not invent facts, values, findings, or conclusions.
+Do not reproduce numerical sequences or repetitive patterns.
+Write one coherent summary of approximately 150–200 words.
+When the summary is complete, stop immediately.
+Begin with: "This study..."
+""".strip()
 
-# Ensure only Question + A-D columns
-MCQA_DF = MCQA_DF[["Question", "A", "B", "C", "D"]]
 
-# --- 3. Lazy Loading for Models ---
+# ============================================================
+# Global model state
+# ============================================================
 _model = None
 _tokenizer = None
-
-# --- Dynamic adapter configuration for agent-chat testing ---
-AGENT_ADAPTER_REPO_ID = "starfriend/WaterScopeAI-Adapters"
+_model_lock = threading.Lock()
 
 
-def download_adapter_subfolder(subfolder: str) -> Path:
-    """Download one adapter folder from the private Hub repository.
-
-    The Space must define an HF_TOKEN secret with read access to
-    starfriend/WaterScopeAI-Adapters. Hugging Face caches unchanged files, so
-    repeated requests do not redownload the same adapter files.
+# ============================================================
+# Load a pristine base model only
+# ============================================================
+def load_base_model():
+    """
+    Load only the original Unsloth-hosted Meta Llama 3.1 8B Instruct model.
+
+    This diagnostic app:
+    - does not import PEFT;
+    - does not load any LoRA adapter;
+    - does not call PeftModel.from_pretrained();
+    - does not call load_adapter();
+    - does not call merge_and_unload().
     """
-    hf_token = os.getenv("waterscopeai")
-
-    if not hf_token:
-        raise RuntimeError(
-            "HF_TOKEN is not configured. Add a secret named HF_TOKEN in "
-            "the Hugging Face Space settings, with read access to "
-            f"{AGENT_ADAPTER_REPO_ID}."
-        )
-
-    local_repo_path = snapshot_download(
-        repo_id=AGENT_ADAPTER_REPO_ID,
-        repo_type="model",
-        revision="main",
-        token=hf_token,
-        allow_patterns=[
-            f"{subfolder}/*",
-            f"{subfolder}/**",
-        ],
-    )
-
-    adapter_path = Path(local_repo_path) / Path(subfolder)
-    config_path = adapter_path / "adapter_config.json"
-
-    if not config_path.is_file():
-        raise FileNotFoundError(
-            "adapter_config.json was not downloaded from "
-            f"{AGENT_ADAPTER_REPO_ID}/{subfolder}. "
-            f"Expected local path: {config_path}"
-        )
-
-    return adapter_path
-
-# User-facing adapter aliases -> repository subfolders
-AGENT_ADAPTER_PATHS = {
-    "da_dpo": "DA-DPO",
-    "da_it": "DA-IT",
-    "dpo": "DPO",
-    "bpd": "DA/BPd",
-    "book": "DA/Book",
-    "paperall": "DA/Paperall",
-    "paperdc": "DA/Paperdc",
-    "tulucore": "DA/Tulucore",
-    "tulufair": "DA/Tulufair",
-    "wikicore": "DA/Wikicore",
-    "wikifair": "DA/Wikifair",
-}
-
-# Unique internal PEFT names prevent conflicts with the existing local
-# adapters named "it" and "dpo", which are used by the unchanged functions.
-AGENT_INTERNAL_NAMES = {
-    alias: f"agent_test_{alias}"
-    for alias in AGENT_ADAPTER_PATHS
-}
-
-_loaded_agent_adapters = set()
-_agent_adapter_lock = threading.Lock()
-
-def load_model_and_tokenizer():
     global _model, _tokenizer
-    
-    print("Initializing model and tokenizer for the first time...")
 
-    if "_model" in globals() and _model is not None:
-        print("Model already initialized.")
+    if _model is not None and _tokenizer is not None:
         return _model, _tokenizer
-        
-    base_model, _tokenizer = FastLanguageModel.from_pretrained(
-        model_name=BASE_MODEL_ID,
-        dtype=None,
-        load_in_4bit=False,
-        device_map={"": 0},  # Use GPU explicitly
-    )
 
-    FastLanguageModel.for_inference(base_model)
+    with _model_lock:
+        if _model is not None and _tokenizer is not None:
+            return _model, _tokenizer
 
-    print("Base model loaded.")
+        print("=" * 70)
+        print("LOADING PRISTINE BASE MODEL ONLY")
+        print(f"Model ID: {BASE_MODEL_ID}")
+        print(f"4-bit loading: {LOAD_IN_4BIT}")
+        print("No PEFT model and no LoRA adapters will be loaded.")
+        print("=" * 70)
 
-    #Load adapters
-    _model = PeftModel.from_pretrained(base_model, IT_ADAPTER_PATH, adapter_name="it")
-    print("DA-IT model loaded.")
-    _model.load_adapter(DPO_ADAPTER_PATH, adapter_name="dpo")
-    print("DA-DPO model loaded.")
+        started = time.time()
 
-    _model.eval()
-    print("All models and adapters are ready!")
-
-# --- 4. Utility Functions ---
-def extract_letter(raw_answer: str) -> str:
-    """Extract predicted option letter from model output"""
-    # Priority 1: Look for explicit phrases like "answer is B"
-    match = re.search(r"(?:answer|option) is\s+([A-H])", raw_answer, re.IGNORECASE)
-    if match:
-        return match.group(1).upper()
+        _model, _tokenizer = FastLanguageModel.from_pretrained(
+            model_name=BASE_MODEL_ID,
+            dtype=None,
+            load_in_4bit=LOAD_IN_4BIT,
+            device_map={"": 0},
+        )
 
-    # Priority 2: Look for formats like "B." or "B)" at the start
-    match = re.search(r"^\s*([A-H])[\.\):]", raw_answer)
-    if match:
-        return match.group(1).upper()
+        FastLanguageModel.for_inference(_model)
+        _model.eval()
 
-    # Priority 3: Look for the first standalone letter in the text
-    match = re.search(r"\b([A-H])\b", raw_answer)
-    if match:
-        return match.group(1).upper()
+        if _tokenizer.pad_token_id is None:
+            _tokenizer.pad_token = _tokenizer.eos_token
 
-    return "N/A"
+        print(f"Loaded in {time.time() - started:.1f} seconds.")
+        print(f"Model class: {type(_model).__name__}")
+        print(f"Tokenizer class: {type(_tokenizer).__name__}")
+        print(f"PEFT configuration present: {hasattr(_model, 'peft_config')}")
+        print("PRISTINE BASE MODEL READY")
+        print("=" * 70)
 
+    return _model, _tokenizer
 
-def clean_repetitions(text: str) -> str:
-    lines = [l.strip() for l in text.strip().splitlines() if l.strip()]
-    if not lines:
-        return ""
 
-    # split into words (keep punctuation as part of word)
-    def tokenize(line):
-        return re.findall(r"\S+", line)
+# ============================================================
+# Message handling
+# ============================================================
+def normalize_history(chat_history):
+    """Convert common Gradio history formats into chat-template messages."""
+    messages = []
 
-    result = tokenize(lines[0])
+    if not chat_history:
+        return messages
 
-    for line in lines[1:]:
-        tokens = tokenize(line)
+    for item in chat_history:
+        if isinstance(item, dict):
+            role = item.get("role")
+            content = item.get("content")
 
-        # find overlap
-        i = 0
-        while i < len(result) and i < len(tokens) and result[i].rstrip(".,!?") == tokens[i].rstrip(".,!?"):
-            i += 1
+            if role in {"user", "assistant"} and content:
+                messages.append(
+                    {"role": role, "content": str(content)}
+                )
 
-        # append only the non-overlapping part
-        result.extend(tokens[i:])
+        elif isinstance(item, (list, tuple)) and len(item) >= 2:
+            user_message, assistant_message = item[0], item[1]
 
-    return " ".join(result)
-    
+            if user_message:
+                messages.append(
+                    {"role": "user", "content": str(user_message)}
+                )
 
-# Global variable to track cancellation
-cancellation_requested = False
+            if assistant_message:
+                messages.append(
+                    {"role": "assistant", "content": str(assistant_message)}
+                )
 
-def run_mcqa_comparison(question, opt_a, opt_b, opt_c, opt_d, opt_e, opt_f, opt_g, opt_h, generate_explanation):
-    global _model, _tokenizer, cancellation_requested
-    
-    # Reset cancellation flag at start
-    cancellation_requested = False
+    return messages
 
-    if _model is None or _tokenizer is None:
-        gr.Info("Initializing models for the first time, this may take a few minutes...")
-        load_model_and_tokenizer()
-
-    options = [opt_a, opt_b, opt_c, opt_d, opt_e, opt_f, opt_g, opt_h]
-    active_options = [opt for opt in options if opt and opt.strip()]
-    if not question or len(active_options) < 2:
-        error_msg = ("Error", "Please enter a question and at least two options.")
-        yield error_msg[0], error_msg[1], error_msg[0], error_msg[1], error_msg[0], error_msg[1]
-        return
-
-    # Check for cancellation
-    if cancellation_requested:
-        gr.Info("Processing cancelled by user")
-        return
-
-    # Build the user input content
-    option_labels = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
-    user_content = [f"Question: {question}\n Candidate options are: "]
-    for i, option_text in enumerate(active_options):
-        user_content.append(f"{option_labels[i]}. {option_text},")
-
-    user_content = " ".join(user_content)
-    
-    if generate_explanation:
-        optional_instruction = "\nProvide the letter of the correct option first. Then give an expert-level explanation in this format: 'The answer is [LETTER]. Because [EXPLANATION].'"
-        max_tokens = 200
-    else:
-        optional_instruction = "\nNo need to explain your answer. "
-        max_tokens = 30
-
-    # System message provides context
-    messages = [
-        {
-            "role": "system",
-            "content": (
-                "You are an expert in water and wastewater treatment and engineering, "
-                "particularly in decarbonization, carbon emissions, and sustainability in the water/wastewater sector. "
-                "Please answer the following multiple choice question. "
-                f"{optional_instruction}"
-            )
-        },
-        {
-            "role": "user",
-            "content": user_content
-        }
-    ]
 
-    # Apply the tokenizer's chat template
-    chat_input = _tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
+def tokenize_safely(chat_text, max_new_tokens):
+    """Tokenize without exceeding the model context window."""
+    model_limit = int(
+        getattr(_model.config, "max_position_embeddings", 131072)
     )
 
-    inputs = _tokenizer(chat_input, return_tensors="pt").to(_model.device)
-
-    def get_prediction(max_tokens_for_generation):
-        # Check for cancellation before generating
-        if cancellation_requested:
-            raise gr.Error("Processing cancelled by user")
-            
-        with torch.no_grad():
-            outputs = _model.generate(
-                **inputs,
-                max_new_tokens=max_tokens_for_generation,
-                eos_token_id=_tokenizer.eos_token_id,
-                pad_token_id=_tokenizer.pad_token_id,
-                temperature=0.000001,
-            )
-        return _tokenizer.decode(outputs[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
-    
-    try:
-        yield "", "Running...", "", "", "", ""
-        
-        # Check for cancellation before processing base model
-        if cancellation_requested:
-            gr.Info("Processing cancelled by user")
-            return
-            
-        with _model.disable_adapter():
-            base_raw = clean_repetitions(get_prediction(max_tokens))
-        base_letter = extract_letter(base_raw)
-        
-        # Check for cancellation before processing IT model
-        if cancellation_requested:
-            gr.Info("Processing cancelled by user")
-            return
-            
-        yield base_letter, base_raw.strip(), "", "Running...", "", ""
-        _model.set_adapter("it")
-        it_raw = clean_repetitions(get_prediction(max_tokens))
-        it_letter = extract_letter(it_raw)
-        
-        # Check for cancellation before processing DPO model
-        if cancellation_requested:
-            gr.Info("Processing cancelled by user")
-            return
-            
-        yield base_letter, base_raw.strip(), it_letter, it_raw.strip(), "", "Running..."
-        _model.set_adapter("dpo")
-        dpo_raw = clean_repetitions(get_prediction(max_tokens))
-        dpo_letter = extract_letter(dpo_raw)
-        
-        yield base_letter, base_raw.strip(), it_letter, it_raw.strip(), dpo_letter, dpo_raw.strip()
-        
-    except gr.Error as e:
-        if "cancelled" in str(e).lower():
-            gr.Info("Processing cancelled by user")
-        else:
-            raise e
-
-
-# Function to handle cancellation
-def cancel_processing():
-    global cancellation_requested
-    cancellation_requested = True
-    return "Cancellation requested"
-
-
-# Safe tokenization wrapper
-def safe_tokenize(chat_input, _tokenizer, _model, max_new_tokens=1600):
-    # 1. Validate input type
-    if not isinstance(chat_input, str) or len(chat_input.strip()) == 0:
-        raise ValueError("chat_input must be a non-empty string")
-
-    # 2. Sanitize weird characters (e.g., emojis, zero-width spaces)
-    clean_input = re.sub(r"[^\x00-\x7F]+", " ", chat_input)
-
-    # 3. Tokenize with truncation to avoid position limit issues
-    max_input_tokens = (
-        _model.config.max_position_embeddings
-        - max_new_tokens
-        - 100
+    # Leave a small safety margin for special tokens.
+    max_input_tokens = max(
+        1,
+        model_limit - int(max_new_tokens) - 32,
     )
 
-    tokens = _tokenizer(
-        clean_input,
+    encoded = _tokenizer(
+        chat_text,
         return_tensors="pt",
         truncation=True,
-        max_length=max_input_tokens
-    )
-
-    # 4. Validate token IDs
-    vocab_size = _model.get_input_embeddings().weight.shape[0]
-    max_id = tokens["input_ids"].max().item()
-    min_id = tokens["input_ids"].min().item()
-
-    print(f"[DEBUG] chat_input: {repr(chat_input)}")
-    print(f"[DEBUG] sanitized_input: {repr(clean_input)}")
-    print(f"[DEBUG] token IDs min: {min_id}, max: {max_id}, vocab size: {vocab_size}")
-
-    if max_id >= vocab_size or min_id < 0:
-        raise ValueError(f"Token IDs out of range: min {min_id}, max {max_id}, vocab size {vocab_size}")
-
-    # 5. Move tokens to model device
-    tokens = {k: v.to(_model.device) for k, v in tokens.items() if isinstance(v, torch.Tensor)}
-    return tokens
-
-    
-# Chat function using DA-DPO model
-def chat_with_model(message, history, max_new_tokens=1600, temperature=0.000001):
-    """
-    Chat with the model using a messages-style prompt (system + history + user).
-    Keeps behavior consistent with the MCQA prompt flow (uses apply_chat_template).
-    """
-    global _model, _tokenizer
-
-    if _model is None or _tokenizer is None:
-        gr.Info("Initializing models for the first time, this may take a few minutes...")
-        load_model_and_tokenizer()
-
-    # Use the DPO adapter (as in your original)
-    _model.set_adapter("dpo")
-    _model.eval()
-
-    system_prompt = (
-                '''
-                You are an expert AI assistant in water and wastewater engineering. 
-                When answering a question:
-                - First, write down all relevant facts or values.  
-                - Next, identify which one is correct based on those facts for factual comparisons or multiple options.
-                - Next, provide a clear description for conceptual definition questions. 
-                - Finally, clearly state your conclusion in this format: 
-                 [Main answer]. [one or two sentences explaining the reasoning]. [appropriate values, equations to support the reasoning].
-                '''
+        max_length=max_input_tokens,
+        add_special_tokens=False,
     )
 
-    # Build messages list (same format used elsewhere)
-    messages = [{"role": "system", "content": system_prompt}]
-
-    if history:
-        for user_msg, bot_msg in history:
-            # skip empty items robustly
-            if user_msg is not None:
-                messages.append({"role": "user", "content": user_msg})
-            if bot_msg is not None:
-                messages.append({"role": "assistant", "content": bot_msg})
-
-    # Current user message
-    messages.append({"role": "user", "content": message})
-
-    # Apply the tokenizer/chat template to get the final chat input string
-    chat_input = _tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-
-    # ---- Use safe tokenization ----
-    try:
-        inputs = safe_tokenize(chat_input, _tokenizer, _model, max_new_tokens=max_new_tokens)
-    except Exception as e:
-        print(f"[ERROR] Failed to tokenize input: {e}")
-        return "Sorry, your input contains unsupported characters. Please rephrase. If error persists, try to clear chat and start the chat again. "
-        
-    # Tokenize (batch of 1) and move tensors to same device as model
-    #inputs = _tokenizer(chat_input, return_tensors="pt").to(_model.device)
-
-    with torch.no_grad():
-        outputs = _model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            eos_token_id=_tokenizer.eos_token_id,
-            pad_token_id=_tokenizer.pad_token_id,
-            do_sample=False,
-            repetition_penalty=1.15,
-            no_repeat_ngram_size=4,
-        )
-
-    # Decode the newly generated tokens only
-    generated_ids = outputs[0][len(inputs["input_ids"][0]):]
-    response = _tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
-    return response
+    encoded = {
+        key: value.to(_model.device)
+        for key, value in encoded.items()
+        if isinstance(value, torch.Tensor)
+    }
 
+    return encoded
 
 
-# Agent chat function supporting the base model, selectable adapters,
-# and a caller-supplied system prompt.
-def agent_chat_with_model(
+# ============================================================
+# Base-only API
+# ============================================================
+def base_chat(
     message,
     chat_history=None,
-    adapter_name="da_it",
     system_prompt="",
-    max_new_tokens=1600,
-    return_diagnostics=False,
+    max_new_tokens=300,
+    return_diagnostics=True,
 ):
     """
-    General question answering and document chat using either the unadapted
-    base model or a selectable LoRA adapter.
-
-    The base model is loaded only once. Each requested adapter is downloaded
-    and loaded on its first use, then reused while the Space remains running.
-
-    API inputs:
-        message: str
-        chat_history: list, optional
-        adapter_name: str
-            One of: base, da_dpo, da_it, dpo, bpd, book, paperall, paperdc,
-            tulucore, tulufair, wikicore, or wikifair.
-        system_prompt: str, optional
-            A true system message supplied by the caller. If empty, the
-            default WaterScope-AI system prompt is used.
-        max_new_tokens: int, optional
-        return_diagnostics: bool, optional
-            If True, prepend model and tokenization diagnostics to the API
-            response so they are visible in the local client output.
-
-    API output:
-        str: assistant response, optionally preceded by diagnostics
+    Generate directly from the pristine base model.
+
+    There is no adapter parameter because this app intentionally contains
+    no PEFT or LoRA code.
     """
-    global _model, _tokenizer, _loaded_agent_adapters
+    global _model, _tokenizer
 
     if _model is None or _tokenizer is None:
-        load_model_and_tokenizer()
+        load_base_model()
 
     if not isinstance(message, str) or not message.strip():
         return "Please provide a non-empty message."
 
-    requested_alias = str(adapter_name or "da_it").strip().lower()
-    available_choices = ["base", *AGENT_ADAPTER_PATHS.keys()]
-
-    if requested_alias not in available_choices:
-        available = ", ".join(available_choices)
-        return (
-            f"Unknown adapter '{requested_alias}'. "
-            f"Available options: {available}."
-        )
+    active_system_prompt = (
+        str(system_prompt).strip()
+        if system_prompt is not None and str(system_prompt).strip()
+        else DEFAULT_SYSTEM_PROMPT
+    )
 
     try:
         max_new_tokens = int(max_new_tokens)
     except (TypeError, ValueError):
-        max_new_tokens = 1600
-
-    max_new_tokens = max(32, min(max_new_tokens, 4096))
-
-    default_system_prompt = """
-You are WaterScope-AI, an expert assistant in water, wastewater,
-environmental engineering, and sustainability.
-
-Answer the user's actual question directly and clearly. You can:
-- explain concepts;
-- answer technical questions;
-- analyze or summarize user-provided text;
-- compare methods, results, or alternatives;
-- discuss research papers and their implications.
-
-When the user provides a document or passage:
-- base document-specific claims only on the supplied content;
-- preserve important quantitative results;
-- do not invent values, citations, standards, organizations, or findings;
-- do not continue or reproduce the source text;
-- do not generate repetitive words, phrases, numbers, or sequences;
-- summarize rather than completing patterns found in the document;
-- clearly distinguish information stated in the document from your interpretation.
-
-For general questions, use your domain knowledge and explain uncertainty when needed.
-Use concise paragraphs unless the user requests a specific structure or level of detail.
-""".strip()
+        max_new_tokens = 300
 
-    active_system_prompt = (
-        str(system_prompt).strip()
-        if system_prompt is not None and str(system_prompt).strip()
-        else default_system_prompt
-    )
+    max_new_tokens = max(16, min(max_new_tokens, 4096))
 
-    messages = [{"role": "system", "content": active_system_prompt}]
-
-    # Accept either Gradio paired history or role/content dictionaries.
-    if chat_history:
-        for item in chat_history:
-            if isinstance(item, dict):
-                role = item.get("role")
-                content = item.get("content")
-                if role in {"user", "assistant"} and content:
-                    messages.append({"role": role, "content": str(content)})
-
-            elif isinstance(item, (list, tuple)) and len(item) >= 2:
-                user_msg, assistant_msg = item[0], item[1]
-                if user_msg:
-                    messages.append({"role": "user", "content": str(user_msg)})
-                if assistant_msg:
-                    messages.append({
-                        "role": "assistant",
-                        "content": str(assistant_msg),
-                    })
-
-    messages.append({"role": "user", "content": message.strip()})
-
-    chat_input = _tokenizer.apply_chat_template(
+    messages = [
+        {"role": "system", "content": active_system_prompt},
+        *normalize_history(chat_history),
+        {"role": "user", "content": message.strip()},
+    ]
+
+    chat_text = _tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True,
     )
 
-    try:
-        inputs = safe_tokenize(
-            chat_input,
-            _tokenizer,
-            _model,
-            max_new_tokens=max_new_tokens,
-        )
-    except Exception as exc:
-        print(f"[ERROR] Agent-chat tokenization failed: {exc}")
-        return (
-            "Sorry, I could not process this input. "
-            "Please shorten the message or clear some chat history and try again."
-        )
+    inputs = tokenize_safely(
+        chat_text=chat_text,
+        max_new_tokens=max_new_tokens,
+    )
 
     diagnostics = {
-        "requested_mode": requested_alias,
+        "test_mode": "pristine_base_only",
         "configured_base_model": BASE_MODEL_ID,
         "model_class": type(_model).__name__,
-        "base_model_class": type(getattr(_model, "base_model", None)).__name__,
-        "active_adapter_before_request": str(
-            getattr(_model, "active_adapter", None)
-        ),
-        "loaded_peft_adapters": sorted(
-            list(getattr(_model, "peft_config", {}).keys())
+        "tokenizer_class": type(_tokenizer).__name__,
+        "peft_config_present": hasattr(_model, "peft_config"),
+        "adapter_methods_present": any(
+            hasattr(_model, method_name)
+            for method_name in (
+                "load_adapter",
+                "set_adapter",
+                "disable_adapter",
+            )
         ),
         "input_tokens": int(inputs["input_ids"].shape[1]),
-        "max_new_tokens": int(max_new_tokens),
+        "max_new_tokens": max_new_tokens,
         "max_position_embeddings": int(
             getattr(_model.config, "max_position_embeddings", -1)
         ),
-        "adapter_disabled_for_generation": requested_alias == "base",
-        "merge_and_unload_used": False,
+        "load_in_4bit": LOAD_IN_4BIT,
+        "do_sample": False,
+        "repetition_penalty": 1.15,
+        "no_repeat_ngram_size": 4,
     }
 
-    print("\n[AGENT CHAT DIAGNOSTICS]")
+    print("\n[BASE-ONLY DIAGNOSTICS]")
     for key, value in diagnostics.items():
         print(f"{key}: {value}")
-    print("[END DIAGNOSTICS]\n")
-
-    # Keep activation and generation under one lock so concurrent requests
-    # cannot switch adapters while another response is being generated.
-    with _agent_adapter_lock:
-        try:
-            if requested_alias == "base":
-                print(
-                    "Agent chat is using the original base-model weights "
-                    "with all PEFT adapters disabled for this generation."
-                )
+    print("[END BASE-ONLY DIAGNOSTICS]\n")
 
-                with _model.disable_adapter():
-                    print(
-                        "Entered _model.disable_adapter() context; "
-                        "LoRA adapters are disabled."
-                    )
-                    _model.eval()
-                    with torch.no_grad():
-                        outputs = _model.generate(
-                            **inputs,
-                            max_new_tokens=max_new_tokens,
-                            do_sample=False,
-                            repetition_penalty=1.15,
-                            no_repeat_ngram_size=4,
-                            eos_token_id=_tokenizer.eos_token_id,
-                            pad_token_id=_tokenizer.pad_token_id,
-                            use_cache=True,
-                        )
-
-            else:
-                subfolder = AGENT_ADAPTER_PATHS[requested_alias]
-                internal_name = AGENT_INTERNAL_NAMES[requested_alias]
-
-                if internal_name not in _loaded_agent_adapters:
-                    print(
-                        f"Downloading agent adapter '{requested_alias}' from "
-                        f"{AGENT_ADAPTER_REPO_ID}/{subfolder}..."
-                    )
-
-                    local_adapter_dir = download_adapter_subfolder(subfolder)
-
-                    print(f"Loading adapter from local path: {local_adapter_dir}")
-                    _model.load_adapter(
-                        str(local_adapter_dir),
-                        adapter_name=internal_name,
-                    )
-                    _loaded_agent_adapters.add(internal_name)
-                    print(
-                        f"Agent adapter '{requested_alias}' loaded successfully."
-                    )
-
-                _model.set_adapter(internal_name)
-                _model.eval()
-                print(f"Agent chat is using adapter: {requested_alias}")
-
-                with torch.no_grad():
-                    outputs = _model.generate(
-                        **inputs,
-                        max_new_tokens=max_new_tokens,
-                        do_sample=False,
-                        repetition_penalty=1.15,
-                        no_repeat_ngram_size=4,
-                        eos_token_id=_tokenizer.eos_token_id,
-                        pad_token_id=_tokenizer.pad_token_id,
-                        use_cache=True,
-                    )
-
-        except Exception as exc:
-            print(
-                f"[ERROR] Failed to use model option "
-                f"'{requested_alias}': {exc}"
-            )
+    with _model_lock:
+        _model.eval()
 
-            if requested_alias == "base":
-                return (
-                    "Could not run the base model with adapters disabled. "
-                    f"Error: {exc}"
-                )
-
-            subfolder = AGENT_ADAPTER_PATHS[requested_alias]
-            return (
-                f"Could not load or use adapter '{requested_alias}'. "
-                f"Repository subfolder: '{subfolder}'. Error: {exc}"
+        with torch.inference_mode():
+            outputs = _model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+                repetition_penalty=1.15,
+                no_repeat_ngram_size=4,
+                eos_token_id=_tokenizer.eos_token_id,
+                pad_token_id=_tokenizer.pad_token_id,
+                use_cache=True,
             )
 
     input_length = inputs["input_ids"].shape[1]
-    generated_ids = outputs[0][input_length:]
+    generated_ids = outputs[0, input_length:]
+
     response = _tokenizer.decode(
         generated_ids,
         skip_special_tokens=True,
@@ -672,8 +262,10 @@ Use concise paragraphs unless the user requests a specific structure or level of
 
     if bool(return_diagnostics):
         diagnostic_text = "\n".join(
-            f"{key}: {value}" for key, value in diagnostics.items()
+            f"{key}: {value}"
+            for key, value in diagnostics.items()
         )
+
         return (
             "=== DIAGNOSTICS ===\n"
             f"{diagnostic_text}\n"
@@ -684,497 +276,79 @@ Use concise paragraphs unless the user requests a specific structure or level of
     return response
 
 
-# Custom CSS for website-like appearance with lighter blue header
-custom_css = """
-.gradio-container {
-    max-width: 1200px !important;
-    margin: 0 auto !important;
-    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
-}
-.header {
-    text-align: center;
-    padding: 20px;
-    background: linear-gradient(135deg, #6eb1ff 0%, #88d3fe 100%);
-    color: white;
-    border-radius: 8px;
-    margin-bottom: 20px;
-}
-.header h1 {
-    margin: 0;
-    font-size: 2.5em;
-    font-weight: 600;
-}
-.header p {
-    margin: 10px 0 0;
-    font-size: 1.2em;
-    opacity: 0.9;
-}
-.section {
-    background: white;
-    padding: 20px;
-    border-radius: 8px;
-    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-    margin-bottom: 20px;
-}
-.nav-bar {
-    margin-bottom: 20px;
-    display: flex;
-    justify-content: center;
-    gap: 10px;
-}
-.footer {
-    text-align: center;
-    padding: 15px;
-    margin-top: 30px;
-    color: #666;
-    font-size: 0.9em;
-    border-top: 1px solid #eee;
-}
-.dataframe-container {
-    margin-top: 20px;
-}
-.model-output {
-    background: #f8f9fa;
-    padding: 15px;
-    border-radius: 8px;
-    border-left: 4px solid #6eb1ff;
-}
-.model-output h4 {
-    margin-top: 0;
-    color: #6eb1ff;
-}
-.option-controls {
-    margin-top: 15px;
-    display: flex;
-    gap: 10px;
-}
-.cancel-btn {
-    background: #f39c12 !important;
-    color: white !important;
-}
-.cancel-btn:hover {
-    background: #e67e22 !important;
-}
-.status-message {
-    padding: 10px;
-    border-radius: 4px;
-    margin: 10px 0;
-}
-.status-info {
-    background-color: #e3f2fd;
-    border-left: 4px solid #2196f3;
-}
-.status-warning {
-    background-color: #fff3e0;
-    border-left: 4px solid #ff9800;
-}
-.status-error {
-    background-color: #ffebee;
-    border-left: 4px solid #f44336;
-}
-.status-success {
-    background-color: #e8f5e9;
-    border-left: 4px solid #4caf50;
-}
-/* Chat specific styles */
-.chat-container {
-    display: flex;
-    flex-direction: column;
-    height: 500px;
-}
-.chat-messages {
-    flex: 1;
-    overflow-y: auto;
-    padding: 15px;
-    background: var(--light);
-    border-radius: 6px;
-    margin-bottom: 15px;
-    display: flex;
-    flex-direction: column;
-    gap: 15px;
-}
-.message {
-    display: flex;
-    max-width: 80%;
-}
-.user-message {
-    align-self: flex-end;
-}
-.bot-message {
-    align-self: flex-start;
-}
-.message-content {
-    padding: 12px 16px;
-    border-radius: 18px;
-    line-height: 1.4;
-}
-.user-message .message-content {
-    background: var(--accent);
-    color: white;
-    border-bottom-right-radius: 4px;
-}
-.bot-message .message-content {
-    background: var(--light-gray);
-    color: var(--dark);
-    border-bottom-left-radius: 4px;
-}
-.chat-input-container {
-    display: flex;
-    gap: 10px;
-}
-.chat-input-container textarea {
-    flex: 1;
-    padding: 12px;
-    border: 1px solid var(--border);
-    border-radius: 6px;
-    resize: vertical;
-    font-family: inherit;
-    font-size: 14px;
-}
+# ============================================================
+# Minimal Gradio Space
+# ============================================================
+with gr.Blocks(title="WaterScope-AI Base Model Diagnostic") as demo:
+    gr.Markdown(
+        """
+# WaterScope-AI: Pristine Base-Model Diagnostic
+
+This temporary Space loads only
+`unsloth/Meta-Llama-3.1-8B-Instruct`.
+
+No PEFT model or LoRA adapter is loaded.
+Use the hidden `/base_chat` API endpoint from your local Python code.
 """
+    )
 
-# --- 5. Gradio UI ---
-with gr.Blocks(theme=gr.themes.Glass(primary_hue="blue"), 
-               title="WaterScope-AI", 
-               fill_width=True,
-               css=custom_css) as demo:
-    
-    # Custom Header with lighter blue
-    with gr.Column(elem_classes="header"):
-        gr.Markdown("WaterScope-AI")
-        gr.Markdown("Domain-Specific Small Language Model for Water Sustainability (Demo)")
-    
-    # Navigation Bar
-    with gr.Row(elem_classes="nav-bar"):
-        gr.Button("Home", variant="secondary", size="sm")
-        gr.Button("About", variant="secondary", size="sm")
-        gr.Button("Documentation", variant="secondary", size="sm")
-        gr.Button("Contact", variant="secondary", size="sm")
-    
-    # Create tabs for different functionalities
-    with gr.Tabs():
-        # MCQA Demo Tab
-        with gr.TabItem("MCQA Demo"):
-            # Status message area
-            status_message = gr.HTML("", elem_classes="status-message")
-            
-            # Main content in a styled section
-            with gr.Column(elem_classes="section"):
-                # State for tracking number of visible options
-                num_options_state = gr.State(4)
-                
-                # Top row with input and output panels
-                with gr.Row():
-                    # Left panel with inputs
-                    with gr.Column(scale=1):
-                        with gr.Group():
-                            question_box = gr.Textbox(label="Question", lines=2, interactive=True)
-
-                            gr.Markdown("#### Options")
-                            
-                            # Create option boxes using a list (like in the working version)
-                            option_boxes = []
-                            for i in range(MAX_OPTIONS):
-                                option_boxes.append(gr.Textbox(
-                                    label=f"Option {chr(ord('A') + i)}", 
-                                    visible=(i < 4), 
-                                    interactive=True
-                                ))
-
-                        with gr.Row():
-                            add_option_btn = gr.Button("Add Option")
-                            clear_btn = gr.Button("Clear")
-                            explanation_checkbox = gr.Checkbox(label="Generate Explanation", value=False)
-
-                        with gr.Row():
-                            run_btn = gr.Button("Run Comparison", variant="primary")
-                            cancel_btn = gr.Button("Cancel", variant="stop", visible=False, elem_classes="cancel-btn")
-
-                    # Right panel with outputs
-                    with gr.Column(scale=1):
-                        gr.Markdown("### Model Outputs")
-                        with gr.Row():
-                            with gr.Column(elem_classes="model-output"):
-                                gr.Markdown("#### Base Model")
-                                base_letter_box = gr.Textbox(label="Predicted Letter", interactive=False)
-                                base_raw_box = gr.Textbox(label="Raw Answer", lines=3, interactive=False)
-                            with gr.Column(elem_classes="model-output"):
-                                gr.Markdown("#### DA-IT Model")
-                                it_letter_box = gr.Textbox(label="Predicted Letter", interactive=False)
-                                it_raw_box = gr.Textbox(label="Raw Answer", lines=3, interactive=False)
-                            with gr.Column(elem_classes="model-output"):
-                                gr.Markdown("#### DA-DPO Model")
-                                dpo_letter_box = gr.Textbox(label="Predicted Letter", interactive=False)
-                                dpo_raw_box = gr.Textbox(label="Raw Answer", lines=3, interactive=False)
-            
-            # Table section with custom styling
-            with gr.Column(elem_classes="section dataframe-container"):
-                gr.Markdown("### Browse 777 MCQAs (click a row to autofill)")
-                mcqa_table = gr.Dataframe(
-                    value=MCQA_DF.values.tolist(),
-                    headers=["Question", "A", "B", "C", "D"],
-                    datatype=["str"]*5,
-                    interactive=False,
-                    wrap=True,
-                    max_height=400
-                )
-        
-        # Chat Tab
-        with gr.TabItem("Chat (for testing)"):
-            with gr.Column(elem_classes="section"):
-                gr.Markdown("### Chat with DA-DPO Model")
-                
-                # Add status indicators for chat
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        gr.Markdown("**System Status:**")
-                        system_status = gr.Textbox(value="Ready", interactive=False)
-                    with gr.Column(scale=1):
-                        gr.Markdown("**API Status:**")
-                        api_status = gr.Textbox(value="Ready", interactive=False)
-                
-                # Chat interface
-                chatbot = gr.Chatbot(
-                    label="Conversation",
-                    elem_classes="chat-messages",
-                    height=400
-                )
-                
-                with gr.Row():
-                    msg = gr.Textbox(
-                        label="Your Message",
-                        placeholder="Type your message here...",
-                        lines=2,
-                        scale=4
-                    )
-                    send_btn = gr.Button("Send", variant="primary", scale=1)
-                
-                clear_chat = gr.Button("Clear Conversation")
-    
-    # Hidden API-only endpoint for base-model and dynamic-adapter testing.
-    # Call with api_name="/agent_chat".
-    agent_api_message = gr.Textbox(visible=False)
-    agent_api_history = gr.JSON(value=[], visible=False)
-    agent_api_adapter = gr.Dropdown(
-        choices=["base", *AGENT_ADAPTER_PATHS.keys()],
-        value="da_it",
+    status = gr.Textbox(
+        value="Base-only diagnostic Space is ready.",
+        label="Status",
+        interactive=False,
+    )
+
+    # Hidden API components
+    api_message = gr.Textbox(
+        label="message",
         visible=False,
     )
-    agent_api_system_prompt = gr.Textbox(
-        value="",
+    api_history = gr.JSON(
+        value=[],
+        label="chat_history",
         visible=False,
     )
-    agent_api_max_tokens = gr.Number(value=1600, precision=0, visible=False)
-    agent_api_return_diagnostics = gr.Checkbox(
-        value=False,
+    api_system_prompt = gr.Textbox(
+        value="",
+        label="system_prompt",
         visible=False,
     )
-    agent_api_output = gr.Textbox(visible=False)
-    agent_api_trigger = gr.Button("Agent Chat API", visible=False)
-
-    agent_api_trigger.click(
-        fn=agent_chat_with_model,
-        inputs=[
-            agent_api_message,
-            agent_api_history,
-            agent_api_adapter,
-            agent_api_system_prompt,
-            agent_api_max_tokens,
-            agent_api_return_diagnostics,
-        ],
-        outputs=agent_api_output,
-        api_name="agent_chat",
+    api_max_tokens = gr.Number(
+        value=300,
+        precision=0,
+        label="max_new_tokens",
+        visible=False,
     )
-
-    # Footer
-    with gr.Column(elem_classes="footer"):
-        gr.Markdown("© 2025 WaterScope-AI | Built with Gradio")
-    
-    # Function to add more options
-    def add_option(current_count):
-        if current_count < MAX_OPTIONS:
-            current_count += 1
-        updates = [gr.update(visible=i < current_count) for i in range(MAX_OPTIONS)]
-        return current_count, *updates
-
-    # Function to clear all inputs and outputs (from working version)
-    def clear_all():
-        """Clears all input and output fields."""
-        option_visibility_updates = [gr.update(visible=(i < 4), value="") for i in range(MAX_OPTIONS)]
-        return (
-            4, # Reset state
-            "", # Clear question
-            *[""]*MAX_OPTIONS, # Clear option values
-            False, # Uncheck the explanation checkbox
-            "", "", "", "", "", "", # Clear all 6 output boxes
-            *option_visibility_updates # Apply visibility and value updates
-        )
-    
-    # Fixed function to load row data
-    def load_row(evt: gr.SelectData):
-        """Load a selected row from the dataframe into the input fields"""
-        if evt.index[0] >= len(MCQA_DF):
-            return ["", ""] + [""] * MAX_OPTIONS
-        
-        row = MCQA_DF.iloc[evt.index[0]]
-        # Return question and first 4 options (A-D), and empty for the rest
-        return_values = [
-            row["Question"] if pd.notna(row["Question"]) else "",
-            row["A"] if pd.notna(row["A"]) else "",
-            row["B"] if pd.notna(row["B"]) else "",
-            row["C"] if pd.notna(row["C"]) else "",
-            row["D"] if pd.notna(row["D"]) else ""
-        ]
-        # Add empty values for any additional options
-        return_values += [""] * (MAX_OPTIONS - 4)
-        return return_values
-
-    # Function to toggle cancel button visibility
-    def toggle_cancel_button():
-        return gr.update(visible=True)
-
-    # Function to hide cancel button
-    def hide_cancel_button():
-        return gr.update(visible=False)
-
-    # Function to update status message
-    def update_status(message, type="info"):
-        if type == "info":
-            cls = "status-info"
-        elif type == "warning":
-            cls = "status-warning"
-        elif type == "error":
-            cls = "status-error"
-        elif type == "success":
-            cls = "status-success"
-        else:
-            cls = "status-info"
-            
-        return f'<div class="status-message {cls}">{message}</div>'
-
-    # Connect the table selection event
-    mcqa_table.select(
-        fn=load_row,
-        inputs=None,
-        outputs=[question_box, *option_boxes]
+    api_return_diagnostics = gr.Checkbox(
+        value=True,
+        label="return_diagnostics",
+        visible=False,
     )
-
-    # Connect the add option button
-    add_option_btn.click(
-        fn=add_option,
-        inputs=[num_options_state],
-        outputs=[num_options_state, *option_boxes]
+    api_output = gr.Textbox(
+        label="response",
+        visible=False,
     )
-    
-    # Define outputs to clear (from working version)
-    outputs_to_clear = [
-        num_options_state, question_box, *option_boxes, explanation_checkbox,
-        base_letter_box, base_raw_box, it_letter_box, it_raw_box, dpo_letter_box, dpo_raw_box,
-        *option_boxes
-    ]
-    
-    # Connect the clear button (from working version)
-    clear_btn.click(
-        fn=clear_all, 
-        inputs=None, 
-        outputs=outputs_to_clear, 
-        queue=False
-    ).then(
-        fn=lambda: update_status("Form cleared", "info"),
-        inputs=None,
-        outputs=[status_message],
-        queue=False
+    api_trigger = gr.Button(
+        "Base Chat API",
+        visible=False,
     )
 
-    # Connect the run button
-    run_event = run_btn.click(
-        fn=lambda: update_status("Initializing processing...", "info"),
-        inputs=None,
-        outputs=[status_message],
-        queue=False
-    ).then(
-        fn=toggle_cancel_button,
-        inputs=None,
-        outputs=[cancel_btn],
-        queue=False
-    ).then(
-        fn=run_mcqa_comparison,
-        inputs=[question_box, *option_boxes, explanation_checkbox],
-        outputs=[base_letter_box, base_raw_box, it_letter_box, it_raw_box, dpo_letter_box, dpo_raw_box]
-    ).then(
-        fn=lambda: update_status("Processing completed successfully", "success"),
-        inputs=None,
-        outputs=[status_message],
-        queue=False
-    ).then(
-        fn=hide_cancel_button,
-        inputs=None,
-        outputs=[cancel_btn],
-        queue=False
+    api_trigger.click(
+        fn=base_chat,
+        inputs=[
+            api_message,
+            api_history,
+            api_system_prompt,
+            api_max_tokens,
+            api_return_diagnostics,
+        ],
+        outputs=api_output,
+        api_name="base_chat",
     )
 
-    # Connect the cancel button
-    cancel_btn.click(
-        fn=cancel_processing,
-        inputs=None,
-        outputs=None,
-        queue=False
-    ).then(
-        fn=lambda: update_status("Processing cancelled by user", "warning"),
-        inputs=None,
-        outputs=[status_message],
-        queue=False
-    ).then(
-        fn=hide_cancel_button,
-        inputs=None,
-        outputs=[cancel_btn],
-        queue=False
-    )
 
-    # Chat functionality
-    def respond(message, chat_history):
-        if not message.strip():
-            return "", chat_history
-        
-        # Update status
-        system_status.value = "Processing..."
-        api_status.value = "Sending request to model"
-        
-        try:
-            # Get response from the model
-            bot_message = chat_with_model(message, chat_history)
-            chat_history.append((message, bot_message))
-            
-            # Update status
-            system_status.value = "Ready"
-            api_status.value = "Response received"
-            
-            return "", chat_history
-        except Exception as e:
-            # Update status with error
-            system_status.value = "Error"
-            api_status.value = f"Error: {str(e)}"
-            
-            # Add error message to chat
-            error_msg = f"Sorry, I encountered an error: {str(e)}"
-            chat_history.append((message, error_msg))
-            return "", chat_history
-
-    # Connect the chat send button
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    send_btn.click(respond, [msg, chatbot], [msg, chatbot])
-    
-    # Connect the clear chat button
-    def clear_chat_func():
-        system_status.value = "Ready"
-        api_status.value = "Ready"
-        return None
-    
-    clear_chat.click(clear_chat_func, None, chatbot, queue=False)
-
-# Pre-load the model when the app starts
-print("Pre-loading models...")
-load_model_and_tokenizer()
-print("Models loaded successfully!")
-
-demo.queue().launch(debug=True)
+print("Pre-loading the pristine base model...")
+load_base_model()
+print("Base-only diagnostic Space is ready.")
+
+demo.queue().launch(debug=True)