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
+from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
 
 # 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", "commentary", "final"]
+REQUIRED_CHANNELS = ["analysis", "final"]
 
 # HF Auth - properly handle multiple token env var names
 HF_TOKEN: Optional[str] = (
@@ -286,6 +286,12 @@ 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],
@@ -310,18 +316,30 @@ def zerogpu_generate(full_prompt,
 
         # Tokenize / prepare inputs
         device = next(model.parameters()).device
-        if HARMONY_AVAILABLE and isinstance(full_prompt, list):
-            input_ids = torch.tensor([full_prompt], dtype=torch.long, device=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)
             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 = enc.to(device)
+            inputs = {k: v.to(device) for k, v in enc.items()}
             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
         # 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
@@ -334,11 +352,10 @@ 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=eos_ids,
             logits_processor=logits_processor,
-            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,
+            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,
         )
         
         # Extract generated tokens only