AbstractPhil commited on
Commit
3d65633
·
1 Parent(s): 23ca4d8
Files changed (1) hide show
  1. app.py +111 -22
app.py CHANGED
@@ -370,7 +370,7 @@ def zerogpu_generate(full_prompt,
370
  max_new_tokens=int(gen_kwargs.get("max_new_tokens", MAX_DEF)),
371
  pad_token_id=model.config.pad_token_id,
372
  eos_token_id=tokenizer.eos_token_id,
373
- bad_words_ids=bad_words_ids,
374
  logits_processor=logits_processor,
375
  repetition_penalty=float(gen_kwargs.get("repetition_penalty", 1.2)),
376
  no_repeat_ngram_size=int(gen_kwargs.get("no_repeat_ngram_size", 8)),
@@ -420,6 +420,58 @@ def zerogpu_generate(full_prompt,
420
  if torch.cuda.is_available():
421
  torch.cuda.empty_cache()
422
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
423
  # -----------------------
424
  # GPU Debug: Harmony Inspector
425
  # -----------------------
@@ -517,29 +569,45 @@ def generate_response(message: str, history: List[List[str]], system_prompt: str
517
  rose_enable: bool, rose_alpha: float, rose_score: Optional[float],
518
  rose_tokens: str, rose_json: str,
519
  show_thinking: bool = False,
 
520
  reasoning_effort: str = "high") -> str:
521
  """
522
  Generate response with proper CoT handling using Harmony format.
523
  """
524
  try:
525
- # Build message list
526
  messages = [{"role": "system", "content": system_prompt or SYSTEM_DEF}]
527
-
528
- # Add history
529
  if history:
530
- for turn in history:
531
- if isinstance(turn, (list, tuple)) and len(turn) >= 2:
532
- user_msg, assistant_msg = turn[0], turn[1]
533
- if user_msg:
534
- messages.append({"role": "user", "content": str(user_msg)})
535
- if assistant_msg:
536
- messages.append({"role": "assistant", "content": str(assistant_msg)})
537
-
538
- # Add current message
539
- messages.append({"role": "user", "content": str(message)})
540
-
541
- # Create Harmony-formatted prompt
542
- if HARMONY_AVAILABLE:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
543
  prompt = create_harmony_prompt(messages, reasoning_effort) # returns token IDs
544
  else:
545
  # Fallback to tokenizer template (string)
@@ -573,7 +641,22 @@ def generate_response(message: str, history: List[List[str]], system_prompt: str
573
  rose_map = None
574
 
575
  # Generate with model
576
- channels = zerogpu_generate(
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
577
  prompt,
578
  {
579
  "do_sample": bool(do_sample),
@@ -641,7 +724,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
641
  lines=2
642
  )
643
 
644
- with gr.Accordion("Generation Settings", open=False):
 
 
 
 
 
 
645
  with gr.Row():
646
  temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature")
647
  top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.01, label="Top-p")
@@ -692,9 +781,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
692
  fn=generate_response,
693
  type="messages",
694
  additional_inputs=[
695
- system_prompt, temperature, top_p, top_k, max_new,
696
- do_sample, seed, rose_enable, rose_alpha, rose_score,
697
- rose_tokens, rose_json, show_thinking, reasoning_effort
698
  ],
699
  title="Chat with Mirel",
700
  description="A chain-of-thought model using Harmony format",
 
370
  max_new_tokens=int(gen_kwargs.get("max_new_tokens", MAX_DEF)),
371
  pad_token_id=model.config.pad_token_id,
372
  eos_token_id=tokenizer.eos_token_id,
373
+
374
  logits_processor=logits_processor,
375
  repetition_penalty=float(gen_kwargs.get("repetition_penalty", 1.2)),
376
  no_repeat_ngram_size=int(gen_kwargs.get("no_repeat_ngram_size", 8)),
 
420
  if torch.cuda.is_available():
421
  torch.cuda.empty_cache()
422
 
423
+ # -----------------------
424
+ # Simple (non-Harmony) GPU path — matches your minimal example
425
+ # -----------------------
426
+ @spaces.GPU(duration=120)
427
+ def zerogpu_generate_simple(prompt_str: str, gen_kwargs: Dict[str, Any], rose_map: Optional[Dict[str, float]], rose_alpha: float, seed: Optional[int]) -> Dict[str, str]:
428
+ """Straight chat_template path. No Harmony tokens. Slices completion from prompt_len.
429
+ Mirrors the minimal HF example and avoids header loops entirely."""
430
+ model = None
431
+ try:
432
+ if seed is not None:
433
+ torch.manual_seed(int(seed))
434
+ model = _load_model_on("auto")
435
+ device = next(model.parameters()).device
436
+
437
+ # Encode prompt string
438
+ enc = tokenizer(prompt_str, return_tensors="pt")
439
+ inputs = {k: v.to(device) for k, v in enc.items()}
440
+ prompt_len = int(inputs["input_ids"].shape[1])
441
+ if "attention_mask" not in inputs:
442
+ inputs["attention_mask"] = torch.ones_like(inputs["input_ids"], dtype=torch.long, device=device)
443
+
444
+ # Optional Rose bias
445
+ logits_processor = None
446
+ if rose_map:
447
+ bias = build_bias_from_tokens(tokenizer, rose_map).to(device)
448
+ logits_processor = [RoseGuidedLogits(bias, float(rose_alpha))]
449
+
450
+ out_ids = model.generate(
451
+ **inputs,
452
+ do_sample=bool(gen_kwargs.get("do_sample", True)),
453
+ temperature=float(gen_kwargs.get("temperature", 0.6)),
454
+ top_p=(float(gen_kwargs.get("top_p")) if gen_kwargs.get("top_p") is not None else None),
455
+ top_k=(int(gen_kwargs.get("top_k")) if gen_kwargs.get("top_k") else None),
456
+ max_new_tokens=int(gen_kwargs.get("max_new_tokens", MAX_DEF)),
457
+ pad_token_id=model.config.pad_token_id,
458
+ logits_processor=logits_processor,
459
+ )
460
+ # Slice generated continuation only
461
+ new_ids = out_ids[0, prompt_len:]
462
+ text = tokenizer.decode(new_ids, skip_special_tokens=True)
463
+ return {"final": text}
464
+ except Exception as e:
465
+ return {"final": f"[Error] {type(e).__name__}: {e}"}
466
+ finally:
467
+ try:
468
+ del model
469
+ except Exception:
470
+ pass
471
+ gc.collect()
472
+ if torch.cuda.is_available():
473
+ torch.cuda.empty_cache()
474
+
475
  # -----------------------
476
  # GPU Debug: Harmony Inspector
477
  # -----------------------
 
569
  rose_enable: bool, rose_alpha: float, rose_score: Optional[float],
570
  rose_tokens: str, rose_json: str,
571
  show_thinking: bool = False,
572
+ simple_mode: bool = True, # NEW: default to simple chat_template path
573
  reasoning_effort: str = "high") -> str:
574
  """
575
  Generate response with proper CoT handling using Harmony format.
576
  """
577
  try:
578
+ # Build messages robustly for Gradio type='messages' or legacy tuple format
579
  messages = [{"role": "system", "content": system_prompt or SYSTEM_DEF}]
580
+
581
+ # Add prior turns
582
  if history:
583
+ if isinstance(history, list) and history and isinstance(history[0], dict):
584
+ # history is already a flat list of {'role','content'} dicts
585
+ for m in history:
586
+ role = m.get("role")
587
+ content = m.get("content", "")
588
+ if role in ("user", "assistant"):
589
+ messages.append({"role": role, "content": str(content)})
590
+ else:
591
+ for turn in history:
592
+ if isinstance(turn, (list, tuple)) and len(turn) >= 2:
593
+ u, a = turn[0], turn[1]
594
+ if u:
595
+ messages.append({"role": "user", "content": str(u)})
596
+ if a:
597
+ messages.append({"role": "assistant", "content": str(a)})
598
+
599
+ # Current user message
600
+ if isinstance(message, dict):
601
+ user_text = message.get("content", "")
602
+ else:
603
+ user_text = str(message)
604
+ messages.append({"role": "user", "content": user_text})
605
+
606
+ # FAST PATH: simple chat_template prompt (recommended)
607
+ if simple_mode:
608
+ prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
609
+ # Harmony path (optional)
610
+ elif HARMONY_AVAILABLE:
611
  prompt = create_harmony_prompt(messages, reasoning_effort) # returns token IDs
612
  else:
613
  # Fallback to tokenizer template (string)
 
641
  rose_map = None
642
 
643
  # Generate with model
644
+ if simple_mode:
645
+ channels = zerogpu_generate_simple(
646
+ prompt,
647
+ {
648
+ "do_sample": bool(do_sample),
649
+ "temperature": float(temperature),
650
+ "top_p": float(top_p) if top_p is not None else None,
651
+ "top_k": int(top_k) if top_k > 0 else None,
652
+ "max_new_tokens": int(max_new_tokens),
653
+ },
654
+ rose_map,
655
+ float(rose_alpha),
656
+ int(seed) if seed is not None else None,
657
+ )
658
+ else:
659
+ channels = zerogpu_generate(
660
  prompt,
661
  {
662
  "do_sample": bool(do_sample),
 
724
  lines=2
725
  )
726
 
727
+ with gr.Accordion("Generation Settings ", open=False):
728
+ # NEW: toggle to bypass Harmony and use plain chat_template like your minimal script
729
+ simple_mode = gr.Checkbox(
730
+ value=True,
731
+ label="Use simple chat_template (no Harmony)",
732
+ info="Matches the minimal HF example; safest path for now"
733
+ )
734
  with gr.Row():
735
  temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature")
736
  top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.01, label="Top-p")
 
781
  fn=generate_response,
782
  type="messages",
783
  additional_inputs=[
784
+ system_prompt, temperature, top_p, top_k, max_new,
785
+ do_sample, seed, rose_enable, rose_alpha, rose_score,
786
+ rose_tokens, rose_json, show_thinking, simple_mode, reasoning_effort
787
  ],
788
  title="Chat with Mirel",
789
  description="A chain-of-thought model using Harmony format",