yes

app.py

@@ -11,7 +11,7 @@ from typing import List, Dict, Optional, Any
 from datetime import datetime
 import gradio as gr
 import spaces  # required for ZeroGPU
-from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Import Harmony components
 try:
@@ -47,7 +47,7 @@ ZEROGPU           = os.getenv("ZEROGPU", os.getenv("ZERO_GPU", "0")) == "1"
 LOAD_4BIT         = os.getenv("LOAD_4BIT", "0") == "1"
 
 # Harmony channels for CoT
-REQUIRED_CHANNELS = ["analysis", "final"]
+REQUIRED_CHANNELS = ["analysis", "commentary", "final"]
 
 # HF Auth - properly handle multiple token env var names
 HF_TOKEN: Optional[str] = (
@@ -138,7 +138,6 @@ def _load_model_on(device_map: Optional[str]) -> AutoModelForCausalLM:
     #        peft_kwargs["subfolder"] = ADAPTER_SUBFOLDER
     #    model = PeftModel.from_pretrained(model, ADAPTER_ID, is_trainable=False, **peft_kwargs)
     
-
     model.eval()
     # Ensure a valid pad_token_id is set; some OSS checkpoints reuse eos as pad
     if getattr(model.config, "pad_token_id", None) is None:
@@ -190,15 +189,7 @@ def create_harmony_prompt(messages: List[Dict[str, str]], reasoning_effort: str
                 )
 
         convo = Conversation.from_messages(harmony_messages)
-        rendered = harmony_encoding.render_conversation_for_completion(convo, Role.ASSISTANT)
-        # Ensure assistant header includes a final channel + message start to avoid 'assistantassistant...' loops
-        try:
-            _tail = tokenizer.decode(list(rendered)[-64:], skip_special_tokens=False)
-            if '<|channel|>final<|message|>' not in _tail:
-                rendered = list(rendered) + tokenizer.encode('<|channel|>final<|message|>', add_special_tokens=False)
-        except Exception:
-            rendered = list(rendered)
-        return rendered
+        return harmony_encoding.render_conversation_for_completion(convo, Role.ASSISTANT)
 
     # Fallback: tokenizer chat template -> string prompt
     if not messages or messages[0].get("role") != "system":
@@ -282,7 +273,7 @@ def build_bias_from_tokens(tokenizer, mapping: Dict[str, float]) -> torch.Tensor
             for t in tid:
                 if isinstance(t, int) and t >= 0:
                     bias[t] += float(w) / max(1, len(tid))
-        elif isinstance(tid, int) and t >= 0:
+        elif isinstance(tid, int) and tid >= 0:
             bias[tid] += float(w)
     return bias
 
@@ -295,12 +286,6 @@ class RoseGuidedLogits(torch.nn.Module):
     def forward(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
         return scores + self.alpha * self.bias_vec.to(scores.device)
 
-class StopOnTokens(StoppingCriteria):
-    def __init__(self, stop_ids: List[int]):
-        self.stop_ids = set(int(s) for s in (stop_ids or []))
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
-        return int(input_ids[0, -1]) in self.stop_ids
-
 @spaces.GPU(duration=120)
 def zerogpu_generate(full_prompt,
                     gen_kwargs: Dict[str, Any],
@@ -325,42 +310,21 @@ def zerogpu_generate(full_prompt,
 
         # Tokenize / prepare inputs
         device = next(model.parameters()).device
-        if HARMONY_AVAILABLE and not isinstance(full_prompt, str):
-            # Accept list/tuple or any iterable of ints from openai_harmony
-            try:
-                token_list = list(full_prompt)
-            except TypeError:
-                token_list = list(getattr(full_prompt, "ids", getattr(full_prompt, "token_ids", [])))
-            if not token_list:
-                raise ValueError("Harmony prompt produced no tokens")
-            input_ids = torch.tensor([token_list], dtype=torch.long, device=device)
+        if HARMONY_AVAILABLE and isinstance(full_prompt, list):
+            input_ids = torch.tensor([full_prompt], dtype=torch.long, device=device)
             attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=device)
             inputs = {"input_ids": input_ids, "attention_mask": attention_mask}
             prompt_len = input_ids.shape[1]
         else:
             enc = tokenizer(full_prompt, return_tensors="pt")
-            inputs = {k: v.to(device) for k, v in enc.items()}
+            inputs = enc.to(device)
             prompt_len = int(inputs["input_ids"].shape[1])
+            # Guarantee attention_mask exists; avoids pad==eos ambiguity warnings
             if "attention_mask" not in inputs:
                 inputs["attention_mask"] = torch.ones_like(inputs["input_ids"], dtype=torch.long, device=device)
-
-        # Prepare stopping
-        sc = None
-        if HARMONY_AVAILABLE and HARMONY_STOP_IDS:
-            sc = StoppingCriteriaList([StopOnTokens(HARMONY_STOP_IDS)])
-
         # Generate
-        # Disallow degenerate header loops
-        bad_words_ids = None
-        try:
-            _B = []
-            for s in ("assistantassistant", "assistant", "<|assistant|>"):
-                ids = tokenizer.encode(s, add_special_tokens=False)
-                if ids:
-                    _B.append(ids)
-            bad_words_ids = _B if _B else None
-        except Exception:
-            pass
+        # Build EOS list: use ONLY Harmony assistant-action stops (per OpenAI docs)
+        eos_ids = HARMONY_STOP_IDS if HARMONY_AVAILABLE else tokenizer.eos_token_id
 
         out_ids = model.generate(
             **inputs,
@@ -370,12 +334,11 @@ def zerogpu_generate(full_prompt,
             top_k=(int(gen_kwargs.get("top_k")) if gen_kwargs.get("top_k") and int(gen_kwargs.get("top_k")) > 0 else None),
             max_new_tokens=int(gen_kwargs.get("max_new_tokens", MAX_DEF)),
             pad_token_id=model.config.pad_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-            bad_words_ids=bad_words_ids,
+            eos_token_id=eos_ids,
             logits_processor=logits_processor,
-            repetition_penalty=float(gen_kwargs.get("repetition_penalty", 1.2)),
-            no_repeat_ngram_size=int(gen_kwargs.get("no_repeat_ngram_size", 8)),
-            stopping_criteria=sc,
+            repetition_penalty=float(gen_kwargs.get("repetition_penalty", 1.1)),
+            no_repeat_ngram_size=int(gen_kwargs.get("no_repeat_ngram_size", 6)),
+            min_new_tokens=1,
         )
         
         # Extract generated tokens only
@@ -421,93 +384,6 @@ def zerogpu_generate(full_prompt,
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-# -----------------------
-# GPU Debug: Harmony Inspector
-# -----------------------
-@spaces.GPU(duration=120)
-def zerogpu_generate_debug(full_prompt, gen_kwargs: Dict[str, Any]) -> Dict[str, Any]:
-    """Minimal GPU path to run a single prompt and return Harmony-parsed output
-    along with short token previews for debugging. Does not use Rose for clarity."""
-    model = None
-    try:
-        model = _load_model_on("auto")
-        device = next(model.parameters()).device
-
-        # Prepare inputs (tokens if Harmony renderer used, else string -> encode)
-        if HARMONY_AVAILABLE and not isinstance(full_prompt, str):
-            token_list = list(full_prompt)
-            if not token_list:
-                raise ValueError("Harmony prompt produced no tokens")
-            input_ids = torch.tensor([token_list], dtype=torch.long, device=device)
-            attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=device)
-            inputs = {"input_ids": input_ids, "attention_mask": attention_mask}
-            prompt_len = input_ids.shape[1]
-        else:
-            enc = tokenizer(full_prompt, return_tensors="pt")
-            inputs = {k: v.to(device) for k, v in enc.items()}
-            if "attention_mask" not in inputs:
-                inputs["attention_mask"] = torch.ones_like(inputs["input_ids"], dtype=torch.long, device=device)
-            prompt_len = int(inputs["input_ids"].shape[1])
-
-        # Harmony stop via stopping criteria
-        sc = StoppingCriteriaList([StopOnTokens(HARMONY_STOP_IDS)]) if (HARMONY_AVAILABLE and HARMONY_STOP_IDS) else None
-
-        out_ids = model.generate(
-            **inputs,
-            do_sample=bool(gen_kwargs.get("do_sample", True)),
-            temperature=float(gen_kwargs.get("temperature", 0.7)),
-            top_p=float(gen_kwargs.get("top_p", 0.9)),
-            top_k=(int(gen_kwargs.get("top_k")) if gen_kwargs.get("top_k") and int(gen_kwargs.get("top_k")) > 0 else None),
-            max_new_tokens=int(gen_kwargs.get("max_new_tokens", MAX_DEF)),
-            pad_token_id=model.config.pad_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-            bad_words_ids=bad_words_ids,
-            stopping_criteria=sc,
-            repetition_penalty=float(gen_kwargs.get("repetition_penalty", 1.15)),
-            no_repeat_ngram_size=int(gen_kwargs.get("no_repeat_ngram_size", 6)),
-        )
-
-        out_list = out_ids[0].tolist()
-        gen_ids = out_list[prompt_len:]
-        # Truncate at first Harmony stop token if present
-        if HARMONY_AVAILABLE and HARMONY_STOP_IDS:
-            for sid in HARMONY_STOP_IDS:
-                if sid in gen_ids:
-                    gen_ids = gen_ids[:gen_ids.index(sid)]
-                    break
-
-        # Parse channels
-        if HARMONY_AVAILABLE:
-            try:
-                channels = parse_harmony_response(gen_ids)
-            except Exception:
-                decoded = tokenizer.decode(gen_ids, skip_special_tokens=False)
-                channels = {"final": extract_final_channel_fallback(decoded), "raw": decoded}
-        else:
-            decoded = tokenizer.decode(gen_ids, skip_special_tokens=False)
-            channels = {"final": extract_final_channel_fallback(decoded), "raw": decoded}
-
-        # Small previews (avoid flooding logs/UI)
-        preview = {
-            "prompt_len": int(prompt_len),
-            "stop_ids": list(HARMONY_STOP_IDS) if HARMONY_AVAILABLE else [],
-            "gen_len": int(len(gen_ids)),
-            "gen_ids_head": gen_ids[:48],
-            "decoded_head": tokenizer.decode(gen_ids[:256], skip_special_tokens=False),
-            "channels": channels,
-        }
-        return preview
-    except Exception as e:
-        return {"error": f"{type(e).__name__}: {e}"}
-    finally:
-        try:
-            del model
-        except Exception:
-            pass
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
 # -----------------------
 # Gradio handlers
 # -----------------------
@@ -605,21 +481,6 @@ def generate_response(message: str, history: List[List[str]], system_prompt: str
     except Exception as e:
         return f"[Error] {type(e).__name__}: {str(e)}"
 
-# -----------------------
-# Extra handler: Harmony Inspector wrapper
-# -----------------------
-
-def harmony_inspect_handler(user_prompt: str, system_prompt: str, reasoning_effort: str):
-    try:
-        msgs = [{"role": "system", "content": system_prompt or SYSTEM_DEF}, {"role": "user", "content": user_prompt or "What is 2+2?"}]
-        prompt = create_harmony_prompt(msgs, reasoning_effort)
-        return zerogpu_generate_debug(
-            prompt,
-            {"do_sample": True, "temperature": 0.7, "top_p": 0.9, "top_k": 0, "max_new_tokens": MAX_DEF}
-        )
-    except Exception as e:
-        return {"error": f"{type(e).__name__}: {e}"}
-
 # -----------------------
 # UI
 # -----------------------
@@ -681,13 +542,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             value=""
         )
 
-    # --- Harmony Inspector UI ---
-    with gr.Accordion("Harmony Inspector", open=False):
-        debug_prompt = gr.Textbox(label="Debug prompt", value="What is 2+2? Reply with just the number.")
-        run_debug = gr.Button("Run Harmony Inspect")
-        debug_out = gr.JSON(label="Parsed Harmony output", value={})
-        run_debug.click(harmony_inspect_handler, inputs=[debug_prompt, system_prompt, reasoning_effort], outputs=[debug_out])
-
     # Chat interface - using only valid parameters
     chat = gr.ChatInterface(
         fn=generate_response,
@@ -697,7 +551,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             do_sample, seed, rose_enable, rose_alpha, rose_score, 
             rose_tokens, rose_json, show_thinking, reasoning_effort
         ],
-        
         title="Chat with Mirel",
         description="A chain-of-thought model using Harmony format",
         examples=[