from pathlib import Path
from datetime import datetime, timezone
from contextlib import redirect_stdout
from io import StringIO
import hashlib, json, platform, time
import torch, transformers, accelerate
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

out=Path('/runs');out.mkdir(exist_ok=True)
notebook=json.loads(Path('/input/llm-playground.ipynb').read_text())
raw=Path('/input/llm-playground.ipynb').read_bytes()
MODEL_ID='Qwen/Qwen3-0.6B'
MODEL_REVISION='c1899de289a04d12100db370d81485cdf75e47ca'
torch.set_num_threads(2)
device='cpu';dtype=torch.float32
print(json.dumps({'phase':'loading','model':MODEL_ID,'revision':MODEL_REVISION}),flush=True)
tokenizer=AutoTokenizer.from_pretrained(MODEL_ID,revision=MODEL_REVISION,token=False,trust_remote_code=False,cache_dir='/runs/model-cache')
model=AutoModelForCausalLM.from_pretrained(MODEL_ID,revision=MODEL_REVISION,token=False,trust_remote_code=False,torch_dtype=dtype,low_cpu_mem_usage=True,use_safetensors=True,cache_dir='/runs/model-cache').to(device)
model.eval()
assert model.config._commit_hash==MODEL_REVISION
function_cell=next(c for c in notebook['cells'] if c.get('id')=='code-6')
exec(''.join(function_cell['source']))
record={'status':'running','startedAt':datetime.now(timezone.utc).isoformat(),'model':MODEL_ID,'modelRevision':MODEL_REVISION,'sourceNotebookSha256':hashlib.sha256(raw).hexdigest(),'generatorCellSha256':hashlib.sha256(''.join(function_cell['source']).encode()).hexdigest(),'environment':{'location':'server CPU (not Google Colab)','architecture':platform.machine(),'python':platform.python_version(),'torch':torch.__version__,'transformers':transformers.__version__,'accelerate':accelerate.__version__,'device':device,'dtype':'float32','threads':torch.get_num_threads()},'generationDefaults':{'enableThinking':False,'doSample':True,'topP':0.8,'topK':20},'cases':[]}
def save():
    temp=out/'results.json.tmp';temp.write_text(json.dumps(record,ensure_ascii=False,indent=2)+'\n');temp.replace(out/'results.json')
def run(case_id,question,temperature=0.7,seed=42,limit=128):
    started=time.perf_counter();result=ask_llm(question,temperature=temperature,seed=seed,max_new_tokens=limit)
    entry={'id':case_id,'question':question,'temperature':temperature,'seed':seed,'maxNewTokens':limit,**result,'seconds':round(time.perf_counter()-started,3)}
    record['cases'].append(entry);save();print(json.dumps({'phase':'generated','id':case_id,'outputTokens':result['output_tokens'],'hitLimit':result['hit_limit']},ensure_ascii=False),flush=True)
    return result
run('first-question','大規模言語モデルとは何ですか。初めて聞く人向けに、日本語で2文以内で説明してください。')
question='雨の日に読みたくなる短編小説の題名を、日本語で3つ考えてください。題名だけを箇条書きにしてください。'
for seed in [42,7]:
    for temperature in [0.3,0.7,1.2]:run(f'temperature-{temperature}-seed-{seed}',question,temperature,seed,96)
run('prompt-brief','APIについて説明して。')
run('prompt-specific','プログラミングを始めたばかりの人向けに、APIを日本語で説明してください。レストランの注文を例にし、説明は3文以内にしてください。')
for days in [7,3]:
    prompt='次の資料だけを使い、質問に日本語で短く答えてください。資料にない情報は補わないでください。回答の最後に根拠の資料番号を示してください。\n\n'+f'【資料】\n資料01：ノートパソコンの貸し出し期間は{days}日です。延長には管理者の承認が必要です。\n\n【質問】\nノートパソコンは何日借りられますか。'
    run(f'grounding-{days}-days',prompt)
sample='AIの仕組みを、ひとつずつ理解する。';ids=tokenizer.encode(sample,add_special_tokens=False)
record['tokenizerExample']={'text':sample,'characters':len(sample),'tokens':len(ids),'ids':ids,'pieces':[tokenizer.decode([i]) for i in ids],'roundTrip':tokenizer.decode(ids)}
assert record['tokenizerExample']['roundTrip']==sample
record['status']='completed';record['finishedAt']=datetime.now(timezone.utc).isoformat();save()
print(json.dumps({'phase':'completed','cases':len(record['cases'])}),flush=True)