# /// script
# requires-python = ">=3.12"
# dependencies = ["marimo==0.24.0", "numpy>=2.0,<3"]
# ///
import marimo
__generated_with = "0.24.0"
app = marimo.App(width="full")

@app.cell
def _():
    """Spark experiments. Small-model extensions, not reproductions of Goodfire's models."""
    import numpy as np
    import json, re, html, sys, os, importlib
    SMOL='HuggingFaceTB/SmolLM2-135M-Instruct'
    SMOL_REV='12fd25f77366fa6b3b4b768ec3050bf629380bac'
    MINI='sentence-transformers/all-MiniLM-L6-v2'
    MINI_REV='1110a243fdf4706b3f48f1d95db1a4f5529b4d41'
    ORIGINAL='Pip brought a paper boat to the creek. Wattle said the water looked too fast. Pip tried anyway, and the boat disappeared around a rock. They walked downstream together. The boat was caught safely in some reeds. Pip thanked Wattle and decided to test the next boat in a puddle.'
    EMOTIONS=['surprise','disgust','anger','happiness','sadness','fear']
    MONTHS=['January','February','March','April','May','June','July','August','September','October','November','December']
    DAYS=['Monday','Tuesday','Wednesday','Thursday','Friday','Saturday','Sunday']

    def sm(x):
        e=np.exp(x-np.max(x,axis=-1,keepdims=True));return e/e.sum(-1,keepdims=True)

    def projection(ref,other=None):
        ref=np.asarray(ref,dtype=float);mean=ref.mean(0);centre=ref-mean
        _,s,v=np.linalg.svd(centre,full_matrices=False);basis=v[:2]
        return (np.asarray(ref if other is None else other)-mean)@basis.T,float((s[:2]**2).sum()/max((s**2).sum(),1e-20))

    def plot(points, labels=None, second=None, axes=('PC1','PC2'), connect=True):
        points=np.asarray(points);second=np.asarray(second) if second is not None else None
        allp=points if second is None else np.vstack([points,second]);scale=max(.001,float(np.abs(allp).max()))
        def xy(p):return 300+245*p[0]/scale,200-150*p[1]/scale
        out=['<svg viewBox="0 0 600 400" role="img" aria-label="Two-dimensional representation plot" style="max-width:100%;background:#f3f7fc"><path d="M40 200H560 M300 35V365" stroke="#bcc8cd"/><text x="525" y="225">PC1</text><text x="310" y="35">PC2</text>']
        for n,arr in enumerate([points] if second is None else [points,second]):
            colour=['#195cca','#a23c76'][n];coords=[xy(p) for p in arr]
            if connect: out.append('<polyline fill="none" stroke="'+colour+'" stroke-width="2" '+('stroke-dasharray="6 4"' if n else '')+' points="'+' '.join(f'{x:.2f},{y:.2f}' for x,y in coords)+'"/>')
            for i,(x,y) in enumerate(coords):
                lab=str(i+1) if labels is None else str(labels[i]);out.append(f'<circle cx="{x:.2f}" cy="{y:.2f}" r="4" fill="{colour}"/><text x="{x+6:.2f}" y="{y-7:.2f}" font-size="12">{html.escape(lab)}</text>')
        return (''.join(out)+'</svg>').replace('>PC1<','>'+html.escape(axes[0])+'<').replace('>PC2<','>'+html.escape(axes[1])+'<')

    def train_addition(epochs=600,seed=7):
        """Train an actual one-hidden-layer network on 80 of 100 mod-10 sums."""
        rng=np.random.default_rng(seed);pairs=np.array([(a,b) for a in range(10) for b in range(10)])
        x=np.concatenate([np.eye(10)[pairs[:,0]],np.eye(10)[pairs[:,1]]],axis=1);y=pairs.sum(1)%10
        split=rng.permutation(100);train,test=split[:80],split[80:]
        w=rng.normal(0,.25,(20,48));b=np.zeros(48);v=rng.normal(0,.15,(48,10));c=np.zeros(10);history=[]
        for epoch in range(int(epochs)):
            h=np.maximum(0,x[train]@w+b);p=sm(h@v+c);d=(p-np.eye(10)[y[train]])/len(train)
            dh=(d@v.T)*(h>0);w-=.6*(x[train].T@dh+1e-4*w);b-=.6*dh.sum(0);v-=.6*(h.T@d+1e-4*v);c-=.6*d.sum(0)
            if epoch%100==0:history.append({'step':epoch,'training_loss':float(-np.log(p[np.arange(80),y[train]]+1e-12).mean())})
        hidden=np.maximum(0,x@w+b);pred=(hidden@v+c).argmax(1)
        # Class means are descriptive summaries, not a causal intervention.
        means=np.array([hidden[y==k].mean(0) for k in range(10)]);points,var=projection(means)
        return {'model':'trained 20 → 48 ReLU → 10 network','seed':seed,'epochs':epochs,'train_indices':train.tolist(),'test_indices':test.tolist(),'train_correct':int((pred[train]==y[train]).sum()),'train_total':80,'test_correct':int((pred[test]==y[test]).sum()),'test_total':20,'pca_variance':var,'rows':[{'a':int(pairs[i,0]),'b':int(pairs[i,1]),'true_mod10_sum':int(y[i]),'prediction':int(pred[i]),'split':'test' if i in test else 'train'} for i in range(100)],'history':history,'svg':plot(points,list(range(10))),'interpretation':'The plotted points are learned hidden-state means grouped by known sum. They need not form a circle. Grouping itself uses labels. A hand-written Fourier circle is a separate mathematical construction.'}

    _NATIVE={}
    def language_model():
        if sys.platform=='emscripten':raise RuntimeError('This generative-model extension needs native Python. Download this notebook and follow the Spark setup guide. The browser activities above still work.')
        if 'smol' not in _NATIVE:
            torch=importlib.import_module('torch')
            transformers=importlib.import_module('transformers')
            AutoTokenizer=transformers.AutoTokenizer;AutoModelForCausalLM=transformers.AutoModelForCausalLM
            tok=AutoTokenizer.from_pretrained(SMOL,revision=SMOL_REV)
            device=os.environ.get('SPARK_DEVICE','cpu')
            if device not in ['cpu','cuda']:raise ValueError('SPARK_DEVICE must be cpu or cuda.')
            if device=='cuda' and not torch.cuda.is_available():raise RuntimeError('CUDA is unavailable. Use CPU or a configured GPU runtime.')
            model=AutoModelForCausalLM.from_pretrained(SMOL,revision=SMOL_REV).eval().to(device)
            _NATIVE['smol']=(tok,model)
        return _NATIVE['smol']

    def hidden_and_logits(text,layer=15,max_tokens=512):
        torch=importlib.import_module('torch')
        tok,model=language_model();ids=tok(text,return_tensors='pt',add_special_tokens=False).input_ids;count=ids.shape[1];ids=ids[:,-max_tokens:].to(model.device)
        with torch.no_grad():out=model(ids,output_hidden_states=True)
        layer=max(0,min(int(layer),len(out.hidden_states)-1))
        return out.hidden_states[layer][0,-1].float().cpu().numpy(),out.logits[0,-1].float().cpu().numpy(),int(count),layer

    def native_calculator(layer=15,steer_to=17,seed=7,custom_domain="months",custom_start=8,custom_offset=6):
        tok,model=language_model();pairs=[(a,b) for a in range(1,9) for b in range(1,9)]
        chat=lambda p:tok.apply_chat_template([{'role':'user','content':p}],tokenize=False,add_generation_prompt=True)
        prompts=[chat(f'Calculate {a} + {b}. Give only the number.') for a,b in pairs];sums=np.array([a+b for a,b in pairs])
        states=[];rows=[]
        for prompt,total in zip(prompts,sums):
            h,logits,_,actual_layer=hidden_and_logits(prompt,layer);states.append(h)
            # Full next-token prediction, not a forced choice among numbers.
            rows.append({'prompt':prompt,'sum':int(total),'next_token':tok.decode([int(logits.argmax())])})
        x=np.array(states);rng=np.random.default_rng(seed);idx=rng.permutation(64);train,test=idx[:48],idx[48:]
        mu=x[train].mean(0);sd=x[train].std(0)+1e-3;z=(x-mu)/sd
        def fourier(n):return np.array([f(2*np.pi*n/p) for p in [2,5,10] for f in [np.cos,np.sin]])
        targets=np.array([fourier(n) for n in sums]);xt=np.column_stack([z[train],np.ones(48)])
        # Dual ridge avoids a large matrix inverse; only training examples fit the probe.
        w=xt.T@np.linalg.solve(xt@xt.T+5*np.eye(48),targets[train]);pred=np.column_stack([z,np.ones(64)])@w
        baseline=targets[train].mean(0);err=float(np.sqrt(((pred[test]-targets[test])**2).mean()));null=float(np.sqrt(((baseline-targets[test])**2).mean()))
        time_cases=[('August + 6 months','What month is 6 months after August? Answer:',14,'February'),('August + 16 months','What month is 16 months after August? Answer:',24,'December'),('Friday + 2 days','What day is 2 days after Friday? Answer:',7,'Sunday'),('13 + 4 hours','What hour is 4 hours after 13:00? Answer:',17,'17'),('23 + 3 hours','What hour is 3 hours after 23:00? Answer:',26,'2')]
        if custom_domain=='months':
            if not 1<=custom_start<=12:raise ValueError('Months use start 1–12.')
            custom=f'What month is {custom_offset} months after {MONTHS[custom_start-1]}?';expected=MONTHS[(custom_start+custom_offset-1)%12]
        elif custom_domain=='weekdays':
            if not 1<=custom_start<=7:raise ValueError('Weekdays use Monday=1 through Sunday=7.')
            custom=f'What day is {custom_offset} days after {DAYS[custom_start-1]}?';expected=DAYS[(custom_start+custom_offset-1)%7]
        else:
            if not 0<=custom_start<=23:raise ValueError('Hours use 0–23.')
            custom=f'What hour is {custom_offset} hours after {custom_start}:00?';expected=str((custom_start+custom_offset)%24)
        time_cases.append(('Your time example',custom,custom_start+custom_offset,expected))
        def generate_short(prompt):
            torch=importlib.import_module('torch')
            ids=tok(prompt,return_tensors='pt',add_special_tokens=False).input_ids.to(model.device)
            with torch.no_grad():out=model.generate(ids,max_new_tokens=40,do_sample=False,pad_token_id=tok.eos_token_id,attention_mask=torch.ones_like(ids))
            return tok.decode(out[0,ids.shape[1]:],skip_special_tokens=True)
        transfer=[]
        for name,prompt,total,expected in time_cases:
            h,logits,_,_=hidden_and_logits(chat(prompt),layer);p=np.append((h-mu)/sd,1)@w
            transfer.append({'task':name,'prompt':prompt,'ordinary_sum':total,'expected_answer':expected,'next_token':tok.decode([int(logits.argmax())]),'generated_answer':generate_short(chat(prompt)),'probe_rmse':float(np.sqrt(((p-fourier(total))**2).mean()))})
        # Intervene on post-block state using inverse of the fitted probe. Includes exact no-op control.
        torch=importlib.import_module('torch')
        prompt=chat('Calculate 7 + 9. Give only the number.');h,base,_,_=hidden_and_logits(prompt,layer)
        current=np.append((h-mu)/sd,1)@w;delta=(fourier(steer_to)-current)@np.linalg.pinv(w[:-1]);delta=delta*sd
        norm=float(np.linalg.norm(delta));cap=.15*float(np.linalg.norm(h));delta=delta*min(1,cap/max(norm,1e-12))
        def intervene(vector):
            def hook(_module,_inputs,output):
                value=output[0] if isinstance(output,tuple) else output
                changed=value.clone();changed[0,-1]+=torch.tensor(vector,dtype=changed.dtype,device=changed.device)
                return (changed,)+output[1:] if isinstance(output,tuple) else changed
            # hidden_states index k is post block k-1; final index includes final norm, so avoid it.
            handle=model.model.layers[actual_layer-1].register_forward_hook(hook)
            try:
                with torch.no_grad():out=model(tok(prompt,return_tensors='pt',add_special_tokens=False).input_ids.to(model.device))
                return out.logits[0,-1].float().cpu().numpy()
            finally:handle.remove()
        if not 1<=actual_layer<len(model.model.layers):raise ValueError('Choose an intermediate layer from 1 to 29; final-normalised states are not valid for this intervention.')
        noop=intervene(np.zeros_like(delta));changed=intervene(delta)
        return {'model':SMOL,'revision':SMOL_REV,'layer':actual_layer,'device':str(model.device),'seed':seed,'training_indices':train.tolist(),'test_indices':test.tolist(),'heldout_probe_rmse':err,'constant_baseline_rmse':null,'next_token_rows':rows,'rows':transfer,'steering':{'prompt':prompt,'requested_sum':steer_to,'delta_norm':float(np.linalg.norm(delta)),'base_next_token':tok.decode([int(base.argmax())]),'steered_next_token':tok.decode([int(changed.argmax())]),'noop_max_logit_error':float(np.abs(base-noop).max()),'mean_absolute_logit_change':float(np.abs(changed-base).mean())},'svg':plot(pred[test,4:6],[int(sums[i]) for i in test],axes=('cos(10)','sin(10)'),connect=False),'plot_axes':'Predicted cos/sin for period 10; Points are held-out prompts, not a trajectory.','interpretation':'Poor probe performance or transfer is evidence against this classroom hypothesis for this model/prompt/layer. Next-token fragments are not scored as full answers. Steering can disrupt computation; a changed answer is not proof of the complete Llama algorithm.'}

    def truth_experiment(steps=200,seed=7):
        # MODEL and forward are injected from the independently trained Brightlab transformer.
        rng=np.random.default_rng(seed);cols=[MODEL['vocab'].index(c) for c in ['red','blue','green','gold']]
        features=[];truth=[];logits=[]
        for tokens,y in MODEL['cases']:
            l,c=forward(tokens,'edited');features.append(c['residual'][-1][-1]);logits.append(l[cols]);truth.append(cols.index(y))
        h=np.array(features);truth=np.array(truth);logits=np.array(logits);test=np.array(MODEL['test_indices']);train=np.array([i for i in range(len(h)) if i not in test])
        mu=h[train].mean(0);sd=h[train].std(0)+1e-3;x=np.column_stack([(h-mu)/sd,np.ones(len(h))])
        # One logistic classifier per proposed colour. It sees frozen internal features, not truth at test time.
        w=rng.normal(0,.01,(x.shape[1],4));target=np.eye(4)[truth]
        for _ in range(300):
            p=1/(1+np.exp(-np.clip(x[train]@w,-30,30)));w-=.08*(x[train].T@(p-target[train])/len(train)+.005*w)
        scores=1/(1+np.exp(-np.clip(x@w,-30,30)));reward=np.column_stack([scores,np.full(len(x),.55)])
        prior=np.column_stack([logits/3,np.max(logits/3,1)-1]);policy=np.zeros((x.shape[1],5));ref=sm(prior)
        for _ in range(int(steps)):
            p=sm(prior[train]+x[train]@policy);adv=reward[train]-.12*(np.log(p+1e-12)-np.log(ref[train]+1e-12));grad=p*(adv-(p*adv).sum(1,keepdims=True));policy+=.2*x[train].T@grad/len(train)
        after=sm(prior+x@policy)
        def metrics(p):
            a=p[test].argmax(1);answered=a!=4;wrong=answered&(a!=truth[test]);return {'questions':len(test),'answered':int(answered.sum()),'wrong_answers':int(wrong.sum()),'correct_answers':int((answered&(a==truth[test])).sum()),'wrong_among_answered':float(wrong.sum()/max(1,answered.sum())),'coverage':float(answered.mean())}
        names=['red','blue','green','gold','I need to check'];rows=[]
        for i in test:
            tokens,_=MODEL['cases'][i];chosen=int(after[i].argmax());rows.append({'case':int(i),'facts':' '.join(MODEL['vocab'][t] for t in tokens),'supported':names[truth[i]],'before':names[int(ref[i].argmax())],'after':names[chosen],'checker_for_after':None if chosen==4 else round(float(scores[i,chosen]),3)})
        return {'model':'Brightlab 10,522-parameter trained transformer, edited checkpoint; frozen hidden-state checker; separately reward-trained 25 × 5 answer policy','seed':seed,'steps':steps,'train_cases':len(train),'reserved_cases':len(test),'checker_accuracy_all_candidate_claims':float(((scores[test]>=.5)==target[test]).mean()),'before':metrics(ref),'after':metrics(after),'rows':rows,'reward':'Learned support score for colour answers; 0.55 for asking to check; KL penalty 0.12 against initial policy. Weights of transformer and checker stay frozen during policy training.','interpretation':'This is a small contextual-bandit policy update on known fictional facts, not full RLFR. A high checker score can be wrong. Ground truth is used to evaluate reserved cases, never as their training reward.'}

    def sentences(text):
        if len(text)>24000:raise ValueError('Select a passage shorter than 24,000 characters.')
        parts=[s.strip() for s in re.findall(r'[^.!?]+[.!?]+[”\"\x27]*|[^.!?]+$',text) if s.strip()]
        if not 3<=len(parts)<=40:raise ValueError('Use 3–40 sentences. Select a shorter passage from longer works.')
        return parts

    def story_inputs(text,ending,context):
        parts=sentences(text);variant=parts[:-1]+[ending.strip() or parts[-1]]
        def inputs(s):return [' '.join(s[:i+1]) if context=='accumulated' else s[i] for i in range(len(s))]
        return parts,variant,inputs(parts)+inputs(variant)

    def story_result(parts,variant,vectors,counts,model,revision,context,limit,pooling,emotion_rows=None):
        n=len(parts);vectors=np.array(vectors);points,var=projection(vectors[:n],vectors);rows=[]
        for i in range(n):
            distance=float(np.linalg.norm(vectors[i]-vectors[n+i]));cos=float(vectors[i]@vectors[n+i]/max(1e-12,np.linalg.norm(vectors[i])*np.linalg.norm(vectors[n+i])))
            rows.append({'sentence':i+1,'original':parts[i],'changed':variant[i],'PC1':float(points[i,0]),'PC2':float(points[i,1]),'changed_PC1':float(points[n+i,0]),'changed_PC2':float(points[n+i,1]),'tokens':counts[i],'changed_tokens':counts[n+i],'truncated':counts[i]>limit or counts[n+i]>limit,'vector_distance':distance,'cosine_similarity':cos})
        return {'model':model,'revision':revision,'context':context,'token_limit':limit,'pooling':pooling,'pca_variance':var,'projection':'PCA fitted to original version only; unchanged basis and mean for variant. Arbitrary component signs. Two axes are not inherently emotions.','rows':rows,'emotion_readouts':emotion_rows,'svg':plot(points[:n],second=points[n:]),'interpretation':'Blue solid = original; pink dashed = replacement ending. A projection can hide differences. Compare full-vector distances. Movement can reflect length, vocabulary, context and truncation as well as narrative changes.'}

    async def browser_story(text,ending,context='accumulated'):
        parts,variant,inputs=story_inputs(text,ending,context)
        globalThis=importlib.import_module('js').globalThis
        # Marimo runs Pyodide in a worker, which has no window. Import the inference module in that worker.
        origin=str(globalThis.location.origin)
        if origin=='null':raise RuntimeError('Open this notebook from the BrightLab site, not a local file URL, or use native mode.')
        module_url=origin+'/spark/model-runtime.js'
        program='(async()=>{if(!globalThis.sparkEmbed){await import('+json.dumps(module_url)+');}return await globalThis.sparkEmbed('+json.dumps(json.dumps(inputs))+');})()'
        data=json.loads(await globalThis.eval(program))
        return story_result(parts,variant,data['vectors'],data['tokenCounts'],data['model'],data['revision'],context,256,'MiniLM mean-pooled encoder; FIRST 256 tokens retained')

    def native_story(text,ending,context='accumulated',layer=15,emotions=False):
        parts,variant,inputs=story_inputs(text,ending,context);vectors=[];counts=[];emotion_rows=[]
        for t in inputs:
            h,_,count,actual=hidden_and_logits(t,layer,512);vectors.append(h);counts.append(count)
        if emotions:
            # Separate prompted behaviour: average token log likelihood of candidate labels, not 0–10 ratings.
            torch=importlib.import_module('torch')
            tok,model=language_model()
            for i,t in enumerate(inputs[:len(parts)]):
                prompt=f'Story: {t}\nThe main emotion in this story is';ids=tok(prompt,return_tensors='pt',add_special_tokens=False).input_ids[:,-480:].to(model.device);scores=[]
                for emotion in EMOTIONS:
                    suffix=tok(' '+emotion,return_tensors='pt',add_special_tokens=False).input_ids.to(model.device);full=torch.cat([ids,suffix],1)
                    with torch.no_grad():logp=model(full).logits[0].log_softmax(-1)
                    values=[float(logp[ids.shape[1]-1+j,token]) for j,token in enumerate(suffix[0])];scores.append(float(np.mean(values)))
                probs=sm(np.array(scores));emotion_rows.append({'sentence':i+1,**{e:float(p) for e,p in zip(EMOTIONS,probs)}})
        return story_result(parts,variant,vectors,counts,SMOL,SMOL_REV,context,512,f'Last-token hidden state at layer {actual}; LAST 512 tokens retained',emotion_rows or None)

    def fetch_book(book):
        # A curated identifier, not an arbitrary URL / private network fetch.
        if str(book) not in ['1524','2591']:raise ValueError('Choose Hamlet (1524) or Grimms’ Fairy Tales (2591).')
        url=f'https://www.gutenberg.org/ebooks/{book}.txt.utf-8'
        if sys.platform=='emscripten':raise RuntimeError('Gutenberg importing uses native Python. In the browser, paste a selected passage or upload a .txt file instead.')
        from urllib.request import urlopen,Request
        with urlopen(Request(url,headers={'User-Agent':'BrightLabSpark classroom research'}),timeout=45) as response:
            data=response.read(3_000_001)
        if len(data)>3_000_000:raise ValueError('This edition exceeds the 3 MB import limit.')
        return {'source':url,'book':str(book),'text':data.decode('utf-8-sig').replace('\r\n','\n').replace('\r','\n'),'note':'Choose and preview a 3–40-sentence passage. A full play is not one model context. The Gutenberg edition includes licence and editorial material.'}

    import marimo as mo
    lesson = {'id': 'y10-calculator', 'year': 10, 'kind': 'calculator', 'title': 'Inside a geometric calculator', 'question': 'How can several circles represent one number — and help a model add?', 'goal': 'Calculate residues, explain how several periodic features disambiguate numbers, and distinguish a pattern in activations from causal evidence.', 'prerequisite': 'Addition, remainders and coordinates; sine and cosine are introduced visually.', 'source': 'https://www.goodfire.com/research/a-geometric-calculator', 'paper': 'https://arxiv.org/html/2605.01148v1', 'minutes': 60, 'basis': 'Goodfire identified a shared addition mechanism in Llama 3.1 8B. Its number features include periods 2, 5 and 10. The paper tests interventions as well as observing patterns.', 'sequence': ['Predict the residue of 17 on each circle. Explain why one circle cannot identify every number.', 'Use the authors’ activation explorer to compare inputs and output for 6 + 8. Record one observation for each period.', 'Contrast those observations with the authors’ steering experiment. What changed inside the model?', 'In Marimo, train a small network and inspect its learned hidden states. Compare it with an openly specified Fourier representation.', 'Extend to a small pretrained language model. Fit a probe on training prompts, then score reserved prompts.'], 'assessment': 'Submit a residue calculation, a comparison from the original demo, and one claim supported by an intervention rather than just a picture.', 'boundary': 'Our circles are exact mathematical illustrations. The embedded demos are Goodfire’s recorded Llama measurements. Our small trained network and open-model experiments are independent extensions, not a replication of the complete paper.'}
    MODEL = None
    forward = None
    engine = {'train_addition':train_addition,'native_calculator':native_calculator,'truth_experiment':truth_experiment,'browser_story':browser_story,'native_story':native_story,'fetch_book':fetch_book,'ORIGINAL':ORIGINAL}
    return (mo, json, sys, engine, lesson,)

@app.cell
def _(mo, lesson):
    mo.md(f"""# BrightLab Spark · Year {lesson['year']}
    ## {lesson['title']}
    **The question:** {lesson['question']}

    **Your goal:** {lesson['goal']}

    [Read Goodfire’s article]({lesson['source']}) · [Research paper]({lesson['paper']}) · [Spark setup guide](https://brightlab-ai-creators.ian347727.chatgpt.site/spark/field-guide)

    **Research boundary:** {lesson['boundary']}

    Work in pairs: predict first, run once, then use a control to challenge your explanation. Results and text stay in this session unless you download them.
    """)
    return 

@app.cell
def _(mo):
    prediction = mo.ui.text_area(label="Before running: what do you predict, and what evidence would change your mind?", full_width=True)
    prediction
    return (prediction,)

@app.cell
def _(mo):
    get_result, set_result = mo.state(None)
    return (get_result, set_result,)

@app.cell
def _(mo, lesson):
    mo.md("""### 1. Does learning addition automatically make a circle?
    Train a real small network on **80 of 100** possible pairs of digits. Reserve the other 20. Its task is the last digit of the sum: 7 + 9 → 6.

    This is a **mod-10 classroom task**, different from the paper’s ordinary addition and calendar remapping. A network may memorise training pairs and fail on new ones. We plot mean hidden states grouped by the known sum; grouping by labels can itself create structure.

    **Your test:** compare 100 and 1,000 steps with the same seed. Count training and reserved successes separately. Does a more attractive plot mean better generalisation?
    """)
    return 

@app.cell
def _(mo):
    train_form = mo.ui.dictionary({"epochs":mo.ui.number(start=100,stop=1500,step=100,value=600,label="Training steps"),"seed":mo.ui.number(start=1,stop=99,value=7,label="Split and initialisation seed")}).form(submit_button_label="Train the small network")
    train_form
    return (train_form,)

@app.cell
def _(train_form, engine, set_result):
    if train_form.value is not None:
        try:
            _r=engine['train_addition'](**train_form.value)
            set_result({"experiment":"learned digit addition","settings":dict(train_form.value),"result":_r})
        except Exception as _e: set_result({"error":str(_e)})
    return 

@app.cell
def _(mo, lesson):
    mo.md("""### 2. Probe a pretrained language model · native Python
    The downloaded notebook can run **SmolLM2-135M-Instruct** on a CPU. It is much smaller than Goodfire’s Llama 3.1 8B.

    We collect 64 last-token hidden states from arithmetic prompts, fit a ridge probe on 48 and score 16 reserved prompts. Targets are cosine and sine at periods 2, 5 and 10. Compare the probe with the training-mean baseline: **lower error is better**. The plot shows held-out period-10 probe predictions, not a discovered perfect circle.

    We transfer the same probe to August + 6, August + 16, Friday + 2, 13:00 + 4 and 23:00 + 3. Finally, we edit an intermediate state toward a chosen sum, cap the change at 15% of the state’s norm, and include an exact no-op control.

    **Year 10:** explain why a probe that fails on reserved prompts is weak evidence. **Year 11:** compare arithmetic and time transfer, then change one layer with the split fixed. A next-token fragment is not a full generated answer.

    The chat template, model revision, split, layer, predictions and intervention measurements are included in your report. No GPU or paid API is required. First use downloads roughly 270 MB of model weights.
    """)
    return 

@app.cell
def _(mo):
    native_form = mo.ui.dictionary({"layer":mo.ui.number(start=1,stop=29,value=15,label="Intermediate layer"),"steer_to":mo.ui.number(start=2,stop=30,value=17,label="Target ordinary sum for intervention"),"seed":mo.ui.number(start=1,stop=99,value=7,label="Train/test split seed"),"custom_domain":mo.ui.dropdown(options=["months","weekdays","hours"],value="months",label="Your time domain"),"custom_start":mo.ui.number(start=0,stop=23,value=8,label="Your start: months 1–12, weekdays 1–7, hours 0–23"),"custom_offset":mo.ui.number(start=0,stop=30,value=6,label="Your added interval")}).form(submit_button_label="Run the pretrained-model experiment (native only)")
    native_form
    return (native_form,)

@app.cell
def _(native_form, engine, set_result):
    if native_form.value is not None:
        try: set_result({"experiment":"SmolLM Fourier probe and transfer","settings":dict(native_form.value),"result":engine['native_calculator'](**native_form.value)})
        except Exception as _e: set_result({"error":str(_e)})
    return 

@app.cell
def _(get_result, mo, json):
    latest = get_result()
    if latest is None:
        display = mo.md("### Your evidence will appear here\nSubmit an experiment above. Model calculations only start when you submit.")
    elif 'error' in latest:
        display = mo.callout(latest['error'],kind="warn")
    else:
        _r=latest['result'];_overview={k:v for k,v in _r.items() if k not in ['svg','rows','next_token_rows','emotion_readouts','history']}
        _items=[mo.md("### Latest completed experiment: "+latest['experiment']),mo.md("Results belong to the submitted settings below. Edits are included only after you submit again."),mo.accordion({'Submitted settings':mo.md('```json\n'+json.dumps(latest['settings'],indent=2)+'\n```')})]
        if 'svg' in _r: _items.append(mo.Html(_r['svg']))
        if 'before' in _r:
            _before=_r['before'];_after=_r['after']
            _items.append(mo.md(f"**Before feedback:** {_before['wrong_answers']} wrong answers out of {_before['answered']} answers, from {_before['questions']} questions. **After feedback:** {_after['wrong_answers']} wrong answers out of {_after['answered']} answers. The policy asked to check on {_after['questions']-_after['answered']} questions. Compare both mistakes and willingness to answer."))
        if 'train_correct' in _r:
            _items.append(mo.md(f"**Training pairs:** {_r['train_correct']} / 80 correct. **Reserved pairs:** {_r['test_correct']} / 20 correct. Did learning the examples teach a rule that works on new pairs?"))
        if 'heldout_probe_rmse' in _r:
            _items.append(mo.md(f"**Reserved probe error:** {_r['heldout_probe_rmse']:.3f}. **Constant baseline error:** {_r['constant_baseline_rmse']:.3f}. Lower is better. A probe that loses to the baseline has not demonstrated reliable number information on these prompts."))
        if 'pca_variance' in _r:
            _items.append(mo.md(f"The two-dimensional map keeps **{100*_r['pca_variance']:.1f}%** of variation in the original reference vectors. Read the numbered rows below to connect points to evidence."))
        _items.append(mo.accordion({'Model, method and full measurements':mo.md('```json\n'+json.dumps(_overview,indent=2)+'\n```')}))
        _items.append(mo.ui.table(_r.get('rows',[]),page_size=12,label="Evidence table"))
        if _r.get('next_token_rows'): _items.append(mo.ui.table(_r['next_token_rows'],page_size=8,label="Actual next-token predictions"))
        if _r.get('emotion_readouts'): _items.extend([mo.md("**Separate prompted readout:** relative likelihoods among six candidate emotion labels, using mean token log likelihood. These are not the paper’s 0–10 ratings or calibrated emotional probabilities."),mo.ui.table(_r['emotion_readouts'])])
        _items.append(mo.download(data=json.dumps(latest,indent=2).encode(),filename="spark-evidence.json",label="Download settings and measured evidence"))
        display=mo.vstack(_items)
    display
    return (latest,)

@app.cell
def _(mo, lesson):
    reflection = mo.ui.text_area(label="After running: state your claim, cite two measurements, name a limitation, and propose a fair next test.",full_width=True)
    mo.vstack([mo.md("### Explain what you found\n"+lesson['assessment']),reflection,mo.md("**Teacher check:** look for a prediction, a controlled comparison, accurate use of measured evidence and a limit on the conclusion. Accept spoken or drawn explanations for younger students. Do not reward a desired result over an honest failed hypothesis.")])
    return (reflection,)

@app.cell
def _(mo, prediction, reflection, latest, json):
    lab_record = {"prediction":prediction.value,"reflection":reflection.value,"latest_experiment":latest}
    mo.download(data=json.dumps(lab_record,indent=2).encode(),filename="spark-lab-record.json",label="Save my prediction, evidence and explanation")
    return 

if __name__ == "__main__":
    app.run()
