Upload 80 files

app.py

@@ -85,15 +85,16 @@ def obliterate(model_choice: str, method_choice: str, progress=gr.Progress()):
     def on_log(msg):
         log_lines.append(msg)
 
-    def on_stage(stage):
+    def on_stage(result):
+        stage_key = result.stage
         icon = {"summon": "\u26a1", "probe": "\u2692\ufe0f", "distill": "\u269b\ufe0f",
-                "excise": "\u2702\ufe0f", "verify": "\u2705", "rebirth": "\u2b50"}.get(stage.key, "\u25b6")
-        log_lines.append(f"\n{icon} {stage.key.upper()} — {stage.description}")
-        progress((list({"summon": 0, "probe": 1, "distill": 2,
-                         "excise": 3, "verify": 4, "rebirth": 5}.values()).index(
-            {"summon": 0, "probe": 1, "distill": 2,
-             "excise": 3, "verify": 4, "rebirth": 5}.get(stage.key, 0)) + 1) / 6,
-                 desc=f"{stage.key.upper()}")
+                "excise": "\u2702\ufe0f", "verify": "\u2705", "rebirth": "\u2b50"}.get(stage_key, "\u25b6")
+        if result.status == "running":
+            log_lines.append(f"\n{icon} {stage_key.upper()} — {result.message}")
+        stage_order = {"summon": 0, "probe": 1, "distill": 2,
+                       "excise": 3, "verify": 4, "rebirth": 5}
+        idx = stage_order.get(stage_key, 0)
+        progress((idx + 1) / 6, desc=f"{stage_key.upper()}")
 
     try:
         from obliteratus.abliterate import AbliterationPipeline
@@ -115,8 +116,8 @@ def obliterate(model_choice: str, method_choice: str, progress=gr.Progress()):
 
         # Keep the model + tokenizer in memory for chat
         with _lock:
-            _state["model"] = pipeline.model
-            _state["tokenizer"] = pipeline.tokenizer
+            _state["model"] = pipeline.handle.model
+            _state["tokenizer"] = pipeline.handle.tokenizer
             _state["status"] = "ready"
 
         log_lines.append("\n" + "=" * 50)
@@ -171,7 +172,7 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     gen_kwargs = {
         **inputs,
-        "max_new_tokens": max_tokens,
+        "max_new_tokens": int(max_tokens),
         "do_sample": temperature > 0,
         "temperature": max(temperature, 0.01),
         "top_p": 0.9,