"""Single-query, single-head toy attention and deterministic whole-document budgeting."""
import math,json
q=[1,0];K=[[1,0],[0,1],[1,1]];V=[[10,0],[0,10],[5,5]]
logits=[sum(x*y for x,y in zip(q,k))/math.sqrt(2) for k in K];z=sum(math.exp(x) for x in logits);w=[math.exp(x)/z for x in logits];out=[sum(w[i]*V[i][j] for i in range(3)) for j in range(2)]
budget=200-20-10-40
orders=[['FACT','A','B'],['A','FACT','B'],['A','B','FACT']]
retained=[{'order':o,'retained':o[:budget//60],'fact_present':'FACT' in o[:budget//60]} for o in orders]
print(json.dumps({'label':'手设矩阵和截断策略；未测试任何LLM位置效应','logits':logits,'attention_weights':w,'output':out,'evidence_budget':budget,'whole_document_budget_cases':retained},ensure_ascii=False,indent=2))
