Upload chat.py with huggingface_hub

chat.py

@@ -0,0 +1,148 @@
+"""Interactive chat with the fine-tuned Gemma4-E2B reasoning model (base +
+LoRA adapter) — streaming output.
+
+Repo: xbruce22/gemma-4-e2b-reasoning-lora
+Base: unsloth/gemma-4-E2B-it
+
+Auto-detects the accelerator (CUDA / Intel XPU / CPU). Loads the base model,
+applies this LoRA adapter, merges it for fast inference, and runs a
+multi-turn chat using the Gemma4 chat template with thinking ON — the model
+emits a <|channel>thought ... <channel|> reasoning block (concise bullets,
+as it was trained) before the final answer.
+
+Install:
+  pip install torch transformers peft
+
+Run:
+  python chat.py
+  python chat.py --repo xbruce22/gemma-4-e2b-reasoning-lora
+  python chat.py --device cpu            # force CPU
+
+In-chat commands:
+  /q       quit          /reset   clear history
+  /raw     toggle raw output (show <|channel>/<channel|>/<turn|> markers)
+  /think   toggle thinking on/off (default ON)
+"""
+import argparse
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor, TextStreamer
+from peft import PeftModel
+
+DEFAULT_REPO = "xbruce22/gemma-4-e2b-reasoning-lora"
+BASE_MODEL = "unsloth/gemma-4-E2B-it"
+
+# Build special-token strings from chr() so this source file never contains
+# literal angle-bracket markers (avoids editor/toolchain mangling).
+CHAN_OPEN = chr(60) + "|channel>thought" + chr(10)
+CHAN_CLOSE = chr(60) + "channel|" + chr(62)
+TURN_END = chr(60) + "turn|" + chr(62)
+THINK = chr(60) + "|think|" + chr(62)
+
+
+def pick_device():
+    if torch.cuda.is_available():
+        return "cuda", torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    if hasattr(torch, "xpu") and torch.xpu.is_available():
+        return "xpu", torch.bfloat16
+    return "cpu", torch.float32
+
+
+def clean_display(text):
+    if CHAN_OPEN in text and CHAN_CLOSE in text:
+        _, _, rest = text.partition(CHAN_OPEN)
+        thought, _, answer = rest.partition(CHAN_CLOSE)
+        return ("\n── thinking ──\n" + thought.strip() +
+                "\n── answer ──\n" + answer.strip())
+    for m in (TURN_END, THINK):
+        text = text.replace(m, "")
+    return text.strip()
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--repo", default=DEFAULT_REPO,
+                    help="HF repo id of the LoRA adapter")
+    ap.add_argument("--device", default=None,
+                    help="force device: cuda | xpu | cpu")
+    args = ap.parse_args()
+
+    device, dtype = pick_device() if args.device is None else (args.device, torch.float32)
+    print(f"device={device}  dtype={dtype}")
+
+    print("Loading processor...")
+    processor = AutoProcessor.from_pretrained(BASE_MODEL)
+    tokenizer = processor.tokenizer
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print(f"Loading base model {BASE_MODEL} ...")
+    base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, dtype=dtype)
+    base = base.to(device)
+    base.config.use_cache = True
+    base.eval()
+
+    print(f"Applying + merging LoRA adapter {args.repo} ...")
+    model = PeftModel.from_pretrained(base, args.repo)
+    model = model.merge_and_unload()
+    model.eval()
+    print("Ready.\n")
+
+    show_raw = False
+    thinking = True
+    messages = [{"role": "system", "content": "You are a helpful, concise assistant."}]
+
+    print(f"Chat ready. /q quit · /reset · /raw · /think   "
+          f"(thinking={'ON' if thinking else 'OFF'})\n")
+
+    while True:
+        try:
+            user = input("you> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nbye."); break
+        if not user:
+            continue
+        if user == "/q":
+            print("bye."); break
+        if user == "/reset":
+            messages = [{"role": "system", "content": "You are a helpful, concise assistant."}]
+            print("(reset)\n"); continue
+        if user == "/raw":
+            show_raw = not show_raw
+            print(f"(display={'raw' if show_raw else 'clean'})\n"); continue
+        if user == "/think":
+            thinking = not thinking
+            print(f"(thinking={'ON' if thinking else 'OFF'})\n"); continue
+
+        messages.append({"role": "user", "content": user})
+        try:
+            text = processor.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True,
+                enable_thinking=thinking)
+        except TypeError:
+            text = processor.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True)
+
+        inputs = processor(text=[text], return_tensors="pt").to(device)
+        for k in list(inputs.keys()):
+            if "token_type" in k or "pixel" in k or "audio" in k:
+                inputs.pop(k)
+
+        print("model> ", end="", flush=True)
+        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        with torch.inference_mode():
+            out_ids = model.generate(
+                **inputs, max_new_tokens=2048, do_sample=True,
+                temperature=1.0, top_p=0.95, top_k=64,
+                pad_token_id=tokenizer.pad_token_id, streamer=streamer)
+        gen_ids = out_ids[0][inputs["input_ids"].shape[1]:]
+        gen_text = tokenizer.decode(gen_ids, skip_special_tokens=False)
+        messages.append({"role": "assistant",
+                         "content": tokenizer.decode(gen_ids, skip_special_tokens=True)})
+        print()
+        if show_raw:
+            print("--- raw ---"); print(gen_text); print("--- end raw ---")
+        print()
+
+
+if __name__ == "__main__":
+    main()