added default configs

app.py

@@ -1,124 +1,244 @@
 import gradio as gr
 from mosaic import Mosaic  # adjust import as needed
 import spaces
-
+import traceback
+from transformers import AutoModelForCausalLM
+import torch
 
 # Maximum number of model textboxes
 MAX_MODELS = 10
+# Cache for loaded models to reuse
+LOADED_MODELS = {}
+
+GPT_CONFIG_MODELS = [
+    "openai-community/gpt2-large",
+    "openai-community/gpt2-medium",
+    "openai-community/gpt2"
+]
 
+Falcon_CONFIG_MODELS = [
+    "tiiuae/Falcon3-10B-Base",
+    "tiiuae/Falcon3-10B-Instruct",
+    "tiiuae/Falcon3-7B-Base",
+    "tiiuae/Falcon3-7B-Instruct"
+]
+
+# Increase model slots
 def update_textboxes(n_visible):
-    """
-    Given the current visible count, increments it by 1 (up to MAX_MODELS)
-    and returns updated visibility settings for all model textboxes.
-    """
     if n_visible < MAX_MODELS:
         n_visible += 1
-    # Create a list of update objects for each textbox: visible if its index is less than n_visible.
-    updates = []
+    tb_updates = [gr.update(visible=(i < n_visible)) for i in range(MAX_MODELS)]
+    btn_updates = [gr.update(visible=(i < n_visible)) for i in range(MAX_MODELS)]
+    status_updates = [gr.update(visible=(i < n_visible)) for i in range(MAX_MODELS)]
+    return (n_visible, *tb_updates, *btn_updates, *status_updates)
+
+# Decrease model slots and clear removed entries
+def remove_textboxes(n_visible):
+    old = n_visible
+    if n_visible > 2:
+        n_visible -= 1
+    new = n_visible
+    # Remove cached models for slots now hidden
+    for idx in range(new, old):
+        LOADED_MODELS.pop(idx+1, None)
+    tb_updates, btn_updates, status_updates = [], [], []
     for i in range(MAX_MODELS):
         if i < n_visible:
-            updates.append(gr.update(visible=True))
+            tb_updates.append(gr.update(visible=True))
+            btn_updates.append(gr.update(visible=True))
+            status_updates.append(gr.update(visible=True))
         else:
-            updates.append(gr.update(visible=False))
-    return n_visible, *updates
+            tb_updates.append(gr.update(visible=False, value=""))
+            btn_updates.append(gr.update(visible=False))
+            status_updates.append(gr.update(visible=False, value="Not loaded"))
+    return (n_visible, *tb_updates, *btn_updates, *status_updates)
 
-# Function to decrease visible textboxes
-def remove_textboxes(n_visible):
+def apply_config1():
     """
-    Given the current visible count, decrements it by 1 (down to 2)
-    and returns updated visibility settings for all model textboxes,
-    clearing the content of any hidden ones.
+    Returns:
+      - new n_visible (number of boxes to show)
+      - new values & visibility for each model textbox
+      - new visibility for each Load button & status box
     """
-    if n_visible > 2:
-        n_visible -= 1
-    updates = []
+    n_vis = len(GPT_CONFIG_MODELS)
+    tb_updates, btn_updates, status_updates = [], [], []
+
     for i in range(MAX_MODELS):
-        if i < n_visible:
-            # keep visible, preserve existing value
-            updates.append(gr.update(visible=True))
+        if i < n_vis:
+            # show this slot, set its value from CONFIG_MODELS
+            tb_updates.append(gr.update(visible=True, value=GPT_CONFIG_MODELS[i]))
+            btn_updates.append(gr.update(visible=True))
+            status_updates.append(gr.update(visible=True, value="Not loaded"))
         else:
-            # hide and clear content
-            updates.append(gr.update(visible=False, value=""))
-    return (n_visible, *updates)
+            # hide all others
+            tb_updates.append(gr.update(visible=False, value=""))
+            btn_updates.append(gr.update(visible=False))
+            status_updates.append(gr.update(visible=False, value="Not loaded"))
+
+    # Return in the same shape as your update_textboxes/remove_textboxes:
+    #   (n_models_state, *all textboxes, *all load buttons, *all status boxes)
+    return (n_vis, *tb_updates, *btn_updates, *status_updates)
 
-@spaces.GPU() 
-def run_scoring(input_text, model1, model2, model3, model4, model5, model6, model7, model8, model9, model10, threshold_choice, custom_threshold):
+def apply_config2():
     """
-    Collect all non-empty model paths, instantiate Mosaic, compute the score,
-    and return a message based on the threshold.
+    Returns:
+      - new n_visible (number of boxes to show)
+      - new values & visibility for each model textbox
+      - new visibility for each Load button & status box
     """
-    model_paths = []
-    for m in [model1, model2, model3, model4, model5, model6, model7, model8, model9, model10]:
-        if m.strip() != "":
-            model_paths.append(m.strip())
-    if len(model_paths) < 2:
-        return "Please enter at least two model paths.", None, None
-    # Choose threshold value
-    if threshold_choice == "default":
-        threshold = 0.0
-    elif threshold_choice == "raid":
-        threshold = 0.23
-    elif threshold_choice == "custom":
-        threshold = custom_threshold
-    else:
-        threshold = 0.0
-    # Instantiate the Mosaic class with the selected model paths.
-    mosaic_instance = Mosaic(model_name_or_paths=model_paths, one_model_mode=False)
-    final_score = mosaic_instance.compute_end_score(input_text)
-    if final_score < threshold:
-        result_message = "This text was probably generated."
-    else:
-        result_message = "This text is likely human-written."
-    return result_message, final_score, threshold
+    n_vis = len(Falcon_CONFIG_MODELS)
+    tb_updates, btn_updates, status_updates = [], [], []
 
-with gr.Blocks() as demo:
+    for i in range(MAX_MODELS):
+        if i < n_vis:
+            # show this slot, set its value from CONFIG_MODELS
+            tb_updates.append(gr.update(visible=True, value=Falcon_CONFIG_MODELS[i]))
+            btn_updates.append(gr.update(visible=True))
+            status_updates.append(gr.update(visible=True, value="Not loaded"))
+        else:
+            # hide all others
+            tb_updates.append(gr.update(visible=False, value=""))
+            btn_updates.append(gr.update(visible=False))
+            status_updates.append(gr.update(visible=False, value="Not loaded"))
+
+    # Return in the same shape as your update_textboxes/remove_textboxes:
+    #   (n_models_state, *all textboxes, *all load buttons, *all status boxes)
+    return (n_vis, *tb_updates, *btn_updates, *status_updates)
+
+# Load a single model and report status
+def load_single_model(model_path, use_bfloat16=True):
+    try:
+        repo = model_path
+        if not repo:
+            return "Error: No path provided"
+        if repo in LOADED_MODELS:
+            return "Loaded"
+        # actual load; may raise
+        model = AutoModelForCausalLM.from_pretrained(
+                    repo,
+                    device_map="auto",
+                    trust_remote_code=True,
+                    torch_dtype=torch.bfloat16 if use_bfloat16 else torch.float32,
+                )
+        model.eval()
+        LOADED_MODELS[repo] = model
+        return "Loaded"
+    except Exception as e:
+        return f"Error loading model: {e}"
+
+# Determine interactive state for Run button
+def check_all_loaded(n_visible, *status_texts):
+    # status_texts are strings: "Loaded" indicates success
+    needed = status_texts[:n_visible]
+    if all(s == "Loaded" for s in needed):
+        return gr.update(interactive=True)
+    return gr.update(interactive=False)
+
+@spaces.GPU()
+def run_scoring(input_text, *args):
+    """
+    args: first MAX_MODELS entries are model paths, followed by threshold_choice and custom_threshold
+    """
+    try:
+        # unpack
+        models = [m.strip() for m in args[:MAX_MODELS] if m.strip()]
+        threshold_choice = args[MAX_MODELS]
+        custom_threshold = args[MAX_MODELS+1]
+        if len(models) < 2:
+            return "Please enter at least two model paths.", None, None
+        threshold = 0.0 if threshold_choice == "default" else custom_threshold
+        mosaic_instance = Mosaic(model_name_or_paths=models, one_model_mode=False, loaded_models=LOADED_MODELS)
+        final_score = mosaic_instance.compute_end_score(input_text)
+        msg = "This text was probably generated." if final_score < threshold else "This text is likely human-written."
+        return msg, final_score, threshold
+    except Exception as e:
+        tb = traceback.format_exc()
+        return f"Error: {e}\n{tb}", None, None
+
+# Build Blocks UI
+demo = gr.Blocks()
+with demo:
     gr.Markdown("# MOSAIC Scoring App")
     with gr.Row():
         input_text = gr.Textbox(lines=10, placeholder="Enter text here...", label="Input Text")
     with gr.Column():
         gr.Markdown("**⚠️ Please make sure all models have the same tokenizer or it won’t work.**")
         gr.Markdown("### Model Paths (at least 2 required)")
-        gr.Markdown("Order matters for model 1 only, the Reference model. Please use the one with the best perplexity on human texts. (The largest LLM if applicable.) GPT2 models are enough to detect easy prompts from chatgpt.")
-        # State to keep track of the number of visible textboxes (starting with 2)
         n_models_state = gr.State(4)
-        # Create 10 textboxes. We'll name them model1, model2, ..., model10.
-        model1 = gr.Textbox(value="Unbabel/TowerBase-13B-v0.1", label="Model 1 Path ", visible=True)
-        model2 = gr.Textbox(value="Unbabel/TowerBase-7B-v0.1", label="Model 2 Path", visible=True)
-        model3 = gr.Textbox(value="meta-llama/Llama-2-7b-chat-hf", label="Model 3 Path", visible=True)
-        model4 = gr.Textbox(value="meta-llama/Llama-2-7b-hf", label="Model 4 Path", visible=True)
-        model5 = gr.Textbox(value="", label="Model 5 Path", visible=False)
-        model6 = gr.Textbox(value="", label="Model 6 Path", visible=False)
-        model7 = gr.Textbox(value="", label="Model 7 Path", visible=False)
-        model8 = gr.Textbox(value="", label="Model 8 Path", visible=False)
-        model9 = gr.Textbox(value="", label="Model 9 Path", visible=False)
-        model10 = gr.Textbox(value="", label="Model 10 Path", visible=False)
-        # Add plus and minus buttons to reveal/hide textboxes.
+        model_inputs, load_buttons, status_boxes = [], [], []
+        for i in range(1, MAX_MODELS+1):
+            with gr.Row():
+                tb = gr.Textbox(label=f"Model {i} Path", value="" if i > 4 else None, visible=(i <= 4))
+                btn = gr.Button("Load", elem_id=f"load_{i}", visible=(i <= 4))
+                status = gr.Textbox(label="Loading status", value="Not loaded", interactive=False, visible=(i <= 4))
+                btn.click(
+                    fn=load_single_model,
+                    inputs=[tb, gr.State(i)],
+                    outputs=status
+                )
+                model_inputs.append(tb)
+                load_buttons.append(btn)
+                status_boxes.append(status)
         with gr.Row():
-            plus_button = gr.Button("+", elem_id="plus_button")
-            minus_button = gr.Button("-", elem_id="minus_button")
-        # Bind click events
-        plus_button.click(
+            plus = gr.Button("Add model slot", elem_id="plus_button")
+            minus = gr.Button("Remove model slot", elem_id="minus_button")
+            config1_btn = gr.Button("Try Basic gpt Configuration")
+            config2_btn = gr.Button("Try Falcon models Configuration")
+        plus.click(
             fn=update_textboxes,
             inputs=n_models_state,
-            outputs=[n_models_state, model1, model2, model3, model4, model5, model6, model7, model8, model9, model10]
+            outputs=[n_models_state, *model_inputs, *load_buttons, *status_boxes]
         )
-        minus_button.click(
+        minus.click(
             fn=remove_textboxes,
             inputs=n_models_state,
-            outputs=[n_models_state, model1, model2, model3, model4, model5, model6, model7, model8, model9, model10]
+            outputs=[n_models_state, *model_inputs, *load_buttons, *status_boxes]
+        )
+        config1_btn.click(
+            fn=apply_config1,
+            inputs=None,                       # no inputs needed
+            outputs=[                          # must match order:
+                n_models_state,                # 1️⃣ the new visible‑count State
+                *model_inputs,                 # 2️⃣ your list of 10 Textboxes
+                *load_buttons,                 # 3️⃣ your list of 10 Load Buttons
+                *status_boxes                  # 4️⃣ your list of 10 Status Textboxes
+            ]
+        )
+        config2_btn.click(
+            fn=apply_config2,
+            inputs=None,                       # no inputs needed
+            outputs=[                          # must match order:
+                n_models_state,                # 1️⃣ the new visible‑count State
+                *model_inputs,                 # 2️⃣ your list of 10 Textboxes
+                *load_buttons,                 # 3️⃣ your list of 10 Load Buttons
+                *status_boxes                  # 4️⃣ your list of 10 Status Textboxes
+            ]
         )
     with gr.Row():
-        threshold_choice = gr.Radio(choices=["default", "raid", "custom"], value="default", label="Threshold Choice")
+        threshold_choice = gr.Radio(choices=["default", "custom"], value="default", label="Threshold Choice")
         custom_threshold = gr.Number(value=0.0, label="Custom Threshold (if 'custom' selected)")
     with gr.Row():
         output_message = gr.Textbox(label="Result Message")
         output_score = gr.Number(label="Final Score")
         output_threshold = gr.Number(label="Threshold Used")
-    run_button = gr.Button("Run Scoring")
+    gr.Markdown("**⚠️ All models need to be loaded before scoring.**")
+    run_button = gr.Button("Run Scoring", interactive=False)
+    # Enable Run button when all statuses reflect "Loaded"
+    for status in status_boxes:
+        status.change(
+            fn=check_all_loaded,
+            inputs=[n_models_state, *status_boxes],
+            outputs=run_button
+        )
+    n_models_state.change(
+        fn=check_all_loaded,
+        inputs=[n_models_state, *status_boxes],
+        outputs=run_button
+    )
     run_button.click(
         fn=run_scoring,
-        inputs=[input_text, model1, model2, model3, model4, model5, model6, model7, model8, model9, model10, threshold_choice, custom_threshold],
+        inputs=[input_text, *model_inputs, threshold_choice, custom_threshold],
         outputs=[output_message, output_score, output_threshold]
     )
-
-demo.launch()
+# Launch
+demo.launch()

mosaic.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Dict
 import numpy as np
 import torch
 import transformers
@@ -49,36 +49,56 @@ def apply_top_p_with_epsilon(logits: torch.Tensor, top_p: float, epsilon: float
     return new_logits
 
 class Mosaic(object):
-    def __init__(self,
-                 model_name_or_paths: List[str],
-                 use_bfloat16: bool = True,
-                 max_token_observed: int = 512,
-                 unigram: Optional[str] = None,
-                 custom_config : Optional[List[bool]] = None,
-                 stupid_mode: bool = False,
-                 one_model_mode: bool = False
-                 ) -> None:
+    def __init__(
+        self,
+        model_name_or_paths: List[str],
+        use_bfloat16: bool = True,
+        max_token_observed: int = 512,
+        unigram: Optional[str] = None,
+        custom_config: Optional[List[bool]] = None,
+        stupid_mode: bool = False,
+        one_model_mode: bool = False,
+        # new optional argument: preloaded models dict
+        loaded_models: Optional[Dict[str, AutoModelForCausalLM]] = None,
+    ) -> None:
+        """
+        If `loaded_models` is provided, re-use any entries matching
+        model_name_or_paths; otherwise load and optionally register
+        into that dict.
+        """
         self.models = []
-        for i, model_name_or_path in enumerate(model_name_or_paths):
-            model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
-                                                         device_map="auto",
-                                                         trust_remote_code=True,
-                                                         torch_dtype=torch.bfloat16 if use_bfloat16
-                                                         else torch.float32
-                                                         )
-            model.eval()  # Set the model to evaluation mode
+        # ensure we have a dict to cache into if passed
+        cache = loaded_models if loaded_models is not None else {}
+
+        for model_name_or_path in model_name_or_paths:
+            # reuse if already loaded
+            if loaded_models is not None and model_name_or_path in cache:
+                model = cache[model_name_or_path]
+            else:
+                print("Reloading a model was necessary, you probably messed up.")
+                # load from pre-trained hub or path
+                model = AutoModelForCausalLM.from_pretrained(
+                    model_name_or_path,
+                    device_map="auto",
+                    trust_remote_code=True,
+                    torch_dtype=torch.bfloat16 if use_bfloat16 else torch.float32,
+                )
+                model.eval()
+                # cache for reuse
+                if loaded_models is not None:
+                    cache[model_name_or_path] = model
             self.models.append(model)
             print(f"Loaded model: {model_name_or_path}")
-            # Print the device map
-            #print(f"Device map for {model_name_or_path}: {model.hf_device_map}")
+
+        # store optional references
+        self.loaded_models = cache
+        self.one_model_mode = one_model_mode
         
         if stupid_mode:
             self.max_iters = 0
         else:
             self.max_iters = 1000
 
-        self.one_model_mode = one_model_mode
-
         self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_paths[-1])
         if not self.tokenizer.pad_token:
             self.tokenizer.pad_token = self.tokenizer.eos_token