205 lines
7.9 KiB
Python
205 lines
7.9 KiB
Python
#!/usr/bin/env python3
|
|
"""HF Jobs workbench test: SakThai Context 7B merged model, 4-bit, T4 GPU."""
|
|
|
|
import json, time, os, sys
|
|
import torch
|
|
|
|
MODEL = "Nanthasit/sakthai-context-7b-merged"
|
|
HF_TOKEN = os.environ.get("HF_TOKEN", "")
|
|
os.environ["HF_TOKEN"] = HF_TOKEN
|
|
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
|
|
|
print("=" * 60)
|
|
print(f"WORKBENCH TEST — SakThai Context 7B")
|
|
print(f"Model: {MODEL}")
|
|
print(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'NONE'}")
|
|
print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB" if torch.cuda.is_available() else "N/A")
|
|
print(f"Time: {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}")
|
|
print("=" * 60)
|
|
sys.stdout.flush()
|
|
|
|
# Load model with 4-bit quantization
|
|
print("\n📥 Loading model in 4-bit...", flush=True)
|
|
t0 = time.time()
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
|
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
bnb_4bit_use_double_quant=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL, token=HF_TOKEN)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL,
|
|
quantization_config=bnb_config,
|
|
device_map="auto",
|
|
torch_dtype=torch.bfloat16,
|
|
token=HF_TOKEN,
|
|
trust_remote_code=True,
|
|
)
|
|
load_time = time.time() - t0
|
|
print(f"✅ Loaded in {load_time:.1f}s on {model.device}", flush=True)
|
|
print(f" VRAM: {torch.cuda.memory_allocated() / 1e9:.2f}GB used", flush=True)
|
|
|
|
# Tests
|
|
tests = [
|
|
{"name": "basic_greeting", "desc": "Say hello",
|
|
"messages": [
|
|
{"role": "system", "content": "You are SakThai, a helpful assistant. Be concise."},
|
|
{"role": "user", "content": "Say hello in one sentence."}
|
|
]},
|
|
{"name": "tool_call_intent", "desc": "Tool-use intent",
|
|
"messages": [
|
|
{"role": "system", "content": "You are SakThai with tools: search(query), read_file(path), run_command(command)."},
|
|
{"role": "user", "content": "Search for the latest AI news"}
|
|
]},
|
|
{"name": "name_recall", "desc": "Remember name across 3 turns",
|
|
"messages": [
|
|
{"role": "system", "content": "You are SakThai."},
|
|
{"role": "user", "content": "My name is Beer."},
|
|
{"role": "assistant", "content": "Nice to meet you, Beer!"},
|
|
{"role": "user", "content": "What's my name?"}
|
|
]},
|
|
{"name": "factual_qa", "desc": "Capital of Japan",
|
|
"messages": [
|
|
{"role": "system", "content": "You are SakThai. Be concise."},
|
|
{"role": "user", "content": "What is the capital of Japan?"}
|
|
]},
|
|
{"name": "json_output", "desc": "Structured JSON",
|
|
"messages": [
|
|
{"role": "system", "content": "You are SakThai. Respond only with valid JSON."},
|
|
{"role": "user", "content": 'List 3 ML frameworks: {"frameworks": ["a","b","c"]}'}
|
|
]},
|
|
{"name": "instruction_following", "desc": "One sentence only",
|
|
"messages": [
|
|
{"role": "system", "content": "You are SakThai. Exactly one sentence."},
|
|
{"role": "user", "content": "Explain what a transformer is."}
|
|
]},
|
|
{"name": "multi_step_reasoning", "desc": "Math reasoning",
|
|
"messages": [
|
|
{"role": "system", "content": "You are SakThai, a helpful assistant."},
|
|
{"role": "user", "content": "If you have 3 apples and give away 1, then buy 5 more, how many do you have? Show your work."}
|
|
]},
|
|
{"name": "context_window", "desc": "Recall from context",
|
|
"messages": [
|
|
{"role": "system", "content": "You are SakThai."},
|
|
{"role": "user", "content": "'Attention Is All You Need' introduced the transformer architecture with multi-head self-attention, positional encodings, and encoder-decoder structure. BERT, GPT, T5 build on it. What year was the paper published?"}
|
|
]},
|
|
]
|
|
|
|
results = []
|
|
for i, test in enumerate(tests):
|
|
print(f"\n{'─'*60}", flush=True)
|
|
print(f"TEST {i+1}: {test['name']} — {test['desc']}", flush=True)
|
|
|
|
try:
|
|
prompt = tokenizer.apply_chat_template(
|
|
test["messages"], tokenize=False, add_generation_prompt=True
|
|
)
|
|
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
|
|
|
t1 = time.time()
|
|
with torch.no_grad():
|
|
outputs = model.generate(
|
|
**inputs,
|
|
max_new_tokens=256,
|
|
temperature=0.1,
|
|
do_sample=True,
|
|
pad_token_id=tokenizer.eos_token_id,
|
|
)
|
|
elapsed = time.time() - t1
|
|
|
|
input_len = inputs.input_ids.shape[1]
|
|
response = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()
|
|
prompt_tokens = input_len
|
|
completion_tokens = outputs.shape[1] - input_len
|
|
|
|
# Quality checks
|
|
checks = []
|
|
if len(response) > 0: checks.append("non_empty")
|
|
if len(response) > 10: checks.append("substantial")
|
|
if test["name"] == "name_recall" and "beer" in response.lower():
|
|
checks.append("name_recall")
|
|
if test["name"] == "factual_qa" and "tokyo" in response.lower():
|
|
checks.append("correct")
|
|
if test["name"] == "context_window" and "2017" in response:
|
|
checks.append("correct_answer")
|
|
if test["name"] == "json_output":
|
|
try:
|
|
json.loads(response)
|
|
checks.append("valid_json")
|
|
except:
|
|
pass
|
|
if test["name"] == "multi_step_reasoning":
|
|
if any(c in response for c in ["7", "seven"]) and ("apple" in response.lower()):
|
|
checks.append("correct_answer")
|
|
|
|
passed = len(response) > 0 # at minimum non-empty
|
|
result = {
|
|
"name": test["name"], "passed": passed,
|
|
"response_preview": response[:200],
|
|
"response_length": len(response),
|
|
"latency_seconds": round(elapsed, 2),
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"checks": checks,
|
|
}
|
|
status = "✅" if passed else "❌"
|
|
print(f" {status} {response[:120]}", flush=True)
|
|
print(f" ⏱ {elapsed:.2f}s | 📝 {prompt_tokens}→{completion_tokens} | ✅ {checks}", flush=True)
|
|
|
|
except Exception as e:
|
|
result = {"name": test["name"], "passed": False, "error": str(e)[:300]}
|
|
print(f" ❌ {e}", flush=True)
|
|
|
|
results.append(result)
|
|
sys.stdout.flush()
|
|
|
|
# Summary
|
|
print(f"\n{'='*60}", flush=True)
|
|
passed = sum(1 for r in results if r.get("passed"))
|
|
total = len(results)
|
|
print(f"📊 7B WORKBENCH SUMMARY: {passed}/{total} passed", flush=True)
|
|
for r in results:
|
|
status = "✅" if r.get("passed") else "❌"
|
|
lat = f"{r.get('latency_seconds',0):.1f}s" if r.get("passed") else " - "
|
|
detail = str(r.get("checks", r.get("error","?")[:60]))
|
|
print(f" {status} {r['name']:<22} ⏱ {lat} {detail}", flush=True)
|
|
|
|
# Save record
|
|
record = {
|
|
"type": "workbench_hf_jobs",
|
|
"model": MODEL,
|
|
"load_time_seconds": round(load_time, 1),
|
|
"device": str(model.device),
|
|
"gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",
|
|
"vram_used_gb": round(torch.cuda.memory_allocated() / 1e9, 2) if torch.cuda.is_available() else 0,
|
|
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
|
"results": results,
|
|
"summary": f"{passed}/{total} passed",
|
|
}
|
|
|
|
record_path = "/tmp/sakthai-7b-workbench-record.json"
|
|
with open(record_path, "w") as f:
|
|
json.dump(record, f, indent=2)
|
|
print(f"\n💾 Saved: {record_path}", flush=True)
|
|
|
|
# Upload record to HF repo
|
|
try:
|
|
from huggingface_hub import HfApi, login
|
|
login(token=HF_TOKEN)
|
|
api = HfApi()
|
|
api.upload_file(
|
|
path_or_fileobj=record_path,
|
|
path_in_repo=f"eval/workbench-7b-{time.strftime('%Y-%m-%d')}.json",
|
|
repo_id="Nanthasit/sakthai-context-7b-merged",
|
|
repo_type="model",
|
|
)
|
|
print(f"📤 Uploaded to HF repo", flush=True)
|
|
except Exception as e:
|
|
print(f"⚠️ Upload failed: {e}", flush=True)
|
|
|
|
print("\n🏁 Done.", flush=True)
|