Upload eval_zero_shot_cpu.py

eval_zero_shot_cpu.py

@@ -1,24 +1,7 @@
 #!/usr/bin/env python3
 """
-Zero-shot eval of Gemma 4 E2B (2B) on CPU — no GPU needed.
-Use this to test the smallest Gemma 4 before any mobile conversion.
-
-REQUIREMENTS:
-  pip install transformers datasets torch huggingface_hub
-
-USAGE:
-  # Quick test (20 samples, ~2-3 min on laptop CPU)
-  python eval_zero_shot_cpu.py --limit 20
-
-  # Full test split (~400 samples, ~30-45 min on CPU)
-  python eval_zero_shot_cpu.py --limit -1
-
-MODEL SIZE:
-  gemma-4-E2B-it = 2B params
-  FP32 on CPU RAM: ~8 GB peak
-  Use --dtype fp16 to halve RAM to ~4 GB if your CPU supports it.
+Zero-shot eval of Gemma 4 E2B (2B) on CPU/GPU.
 """
-
 import argparse
 import json
 import time
@@ -44,27 +27,22 @@ def parse_args():
     p.add_argument("--model", default="google/gemma-4-E2B-it")
     p.add_argument("--dataset", default="BothBosu/scam-dialogue")
     p.add_argument("--split", default="test")
-    p.add_argument("--limit", type=int, default=20,
-                   help="Max rows (-1 = all). Default 20 for quick CPU test.")
-    p.add_argument("--dtype", default="fp32", choices=["fp16","fp32"],
-                   help="fp16 = half RAM (~4 GB), fp32 = ~8 GB")
-    p.add_argument("--out", default="results_zero_shot_cpu.json")
+    p.add_argument("--limit", type=int, default=50)
+    p.add_argument("--dtype", default="fp16", choices=["fp16","fp32"])
+    p.add_argument("--out", default="results_zero_shot.json")
     return p.parse_args()
 
 
-def load_model_cpu(model_id: str, dtype: str):
+def load_model(model_id: str, dtype: str):
     torch_dtype = torch.float16 if dtype == "fp16" else torch.float32
-    print(f"Loading {model_id} on CPU (dtype={dtype}) …")
-    print(f"  Expected RAM: ~{4 if dtype == 'fp16' else 8} GB")
-    print(f"  If OOM: close browser tabs, reduce --limit, or use fp16.\n")
-
+    print(f"Loading {model_id} (dtype={dtype}) …")
+    
     model = Gemma4ForConditionalGeneration.from_pretrained(
         model_id,
         torch_dtype=torch_dtype,
-        device_map=None,          # force CPU
+        device_map="auto",
         low_cpu_mem_usage=True,
     )
-    model = model.to("cpu")
     processor = AutoProcessor.from_pretrained(model_id)
     model.eval()
     return model, processor
@@ -80,7 +58,7 @@ def classify(model, processor, transcript: str) -> str:
         messages, tokenize=True, return_dict=True,
         return_tensors="pt", add_generation_prompt=True,
     )
-    inputs = {k: v.to("cpu") for k, v in inputs.items()}
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
     gen_ids = model.generate(
         **inputs, max_new_tokens=5, do_sample=False,
@@ -115,7 +93,7 @@ def compute_metrics(items):
 
 def main():
     args = parse_args()
-    model, processor = load_model_cpu(args.model, args.dtype)
+    model, processor = load_model(args.model, args.dtype)
     ds = load_dataset(args.dataset, split=args.split)
     n = len(ds) if args.limit < 0 else min(args.limit, len(ds))
 
@@ -138,10 +116,10 @@ def main():
     metrics["throughput"] = n / elapsed
 
     print("\n" + "=" * 60)
-    print("CPU ZERO-SHOT REPORT")
+    print("ZERO-SHOT EVALUATION REPORT")
     print("=" * 60)
     print(f"Model      : {args.model} (2B params)")
-    print(f"Device     : CPU")
+    print(f"Device     : {next(model.parameters()).device}")
     print(f"Samples    : {n}")
     print(f"Time       : {elapsed:.1f}s ({metrics['throughput']:.2f} ex/s)")
     print(f"Accuracy   : {metrics['accuracy']:.2%}")
@@ -158,8 +136,7 @@ def main():
 
     acc, f1 = metrics["accuracy"], metrics["f1_scam"]
     if acc >= 0.90 and f1 >= 0.85:
-        print("\n✅ PASS — 2B base model is accurate enough.")
-        print("   Next: convert to LiteRT 4-bit for phone (~1.5 GB RAM).")
+        print("\n✅ PASS — 2B base model is accurate enough for phone deployment.")
     elif acc >= 0.75 and f1 >= 0.70:
         print("\n⚠️ MARGINAL — Fine-tune with Unsloth, then LiteRT-convert.")
     else: