okay it works better now

app.py

@@ -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 = ["thinking", "analysis", "final"]
+REQUIRED_CHANNELS = ["analysis", "commentary", "final"]
 
 # HF Auth - properly handle multiple token env var names
 HF_TOKEN: Optional[str] = (
@@ -84,6 +84,9 @@ if HARMONY_AVAILABLE:
 else:
     harmony_encoding = None
 
+# Stop tokens per Harmony spec: <|return|> (200002), <|call|> (200012)
+HARMONY_STOP_IDS = [200002, 200012]
+
 # Tokenizer is lightweight; load once
 try:
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, token=HF_TOKEN)
@@ -144,7 +147,7 @@ def _load_model_on(device_map: Optional[str]) -> AutoModelForCausalLM:
 # Harmony formatting
 # -----------------------
 
-def create_harmony_prompt(messages: List[Dict[str, str]], reasoning_effort: str = "high") -> str:
+def create_harmony_prompt(messages: List[Dict[str, str]], reasoning_effort: str = "high") -> Any:
     """Create a proper Harmony-formatted prompt using openai_harmony."""
     if not HARMONY_AVAILABLE:
         # Fallback to tokenizer's chat template
@@ -189,9 +192,7 @@ def create_harmony_prompt(messages: List[Dict[str, str]], reasoning_effort: str
     # Create conversation and render
     convo = Conversation.from_messages(harmony_messages)
     tokens = harmony_encoding.render_conversation_for_completion(convo, Role.ASSISTANT)
-    
-    # Convert tokens back to text for the model
-    return tokenizer.decode(tokens)
+    return tokens  # pass tokens directly to the model to avoid decode/re-encode drift
 
 def parse_harmony_response(tokens: List[int]) -> Dict[str, str]:
     """Parse response tokens using Harmony format to extract channels."""
@@ -268,7 +269,7 @@ class RoseGuidedLogits(torch.nn.Module):
         return scores + self.alpha * self.bias_vec.to(scores.device)
 
 @spaces.GPU(duration=120)
-def zerogpu_generate(full_prompt: str,
+def zerogpu_generate(full_prompt,
                     gen_kwargs: Dict[str, Any],
                     rose_map: Optional[Dict[str, float]],
                     rose_alpha: float,
@@ -289,8 +290,16 @@ def zerogpu_generate(full_prompt: str,
             eff_alpha = float(rose_alpha) * (float(rose_score) if rose_score is not None else 1.0)
             logits_processor = [RoseGuidedLogits(bias, eff_alpha)]
 
-        # Tokenize input
-        inputs = tokenizer(full_prompt, return_tensors="pt").to(next(model.parameters()).device)
+        # 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)
+            inputs = {"input_ids": input_ids}
+            prompt_len = input_ids.shape[1]
+        else:
+            enc = tokenizer(full_prompt, return_tensors="pt")
+            inputs = enc.to(device)
+            prompt_len = int(inputs["input_ids"].shape[1])
         
         # Generate
         out_ids = model.generate(
@@ -301,21 +310,33 @@ def zerogpu_generate(full_prompt: str,
             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=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
+            eos_token_id=(HARMONY_STOP_IDS if HARMONY_AVAILABLE else tokenizer.eos_token_id),
             logits_processor=logits_processor,
         )
         
         # Extract generated tokens only
-        prompt_len = int(inputs["input_ids"].shape[1])
-        gen_ids = out_ids[0][prompt_len:].tolist()
+        out_list = out_ids[0].tolist()
+        gen_ids = out_list[prompt_len:]
+        # Truncate at first Harmony stop token if present
+        if HARMONY_AVAILABLE:
+            for sid in HARMONY_STOP_IDS:
+                if sid in gen_ids:
+                    gen_ids = gen_ids[:gen_ids.index(sid)]
+                    break
         
         # Parse response with Harmony
         if HARMONY_AVAILABLE:
-            channels = parse_harmony_response(gen_ids)
+            try:
+                channels = parse_harmony_response(gen_ids)
+            except Exception:
+                # Fallback to text parsing if Harmony parser fails
+                decoded = tokenizer.decode(gen_ids, skip_special_tokens=False)
+                channels = {
+                    "final": extract_final_channel_fallback(decoded),
+                    "raw": decoded
+                }
         else:
             # Fallback
-            decoded = tokenizer.decode(gen_ids, skip_special_tokens=False)
-            channels = {
                 "final": extract_final_channel_fallback(decoded),
                 "raw": decoded
             }
@@ -367,9 +388,9 @@ def generate_response(message: str, history: List[List[str]], system_prompt: str
         
         # Create Harmony-formatted prompt
         if HARMONY_AVAILABLE:
-            prompt = create_harmony_prompt(messages, reasoning_effort)
+            prompt = create_harmony_prompt(messages, reasoning_effort)  # returns token IDs
         else:
-            # Fallback to tokenizer template
+            # Fallback to tokenizer template (string)
             prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
 
         # Build Rose map if enabled

requirements.txt

@@ -4,4 +4,5 @@ accelerate>=0.33.0
 peft>=0.11.0
 torch>=2.4.0   # ZeroGPU-supported (2.3.x is NOT supported)
 bitsandbytes>=0.43.1
-openai_harmony
+openai_harmony
+gradio>=5.42.0