File size: 5,506 Bytes
e70aa91 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """Bounded tensor-backed file state, not a transformer-compiled filesystem.
No SQLite. UTF-8 bytes live in the model session's tensors. The controller uses
explicit tensor indexing; it is deliberately not advertised as learned logic.
"""
from dataclasses import dataclass
import torch
from torch import nn
@dataclass
class MemoryState:
names: torch.Tensor
name_length: torch.Tensor
data: torch.Tensor
length: torch.Tensor
version: torch.Tensor
deleted: torch.Tensor
clock: torch.Tensor
def clone(self):
return MemoryState(**{k:v.clone() for k,v in vars(self).items()})
class TensorFiles(nn.Module):
def __init__(self, slots=64, versions=8, name_bytes=128, file_bytes=8192):
super().__init__()
self.spec = dict(slots=slots,versions=versions,name_bytes=name_bytes,file_bytes=file_bytes)
self.register_buffer("anchor",torch.zeros(0),persistent=False)
def empty(self):
s=self.spec; dev=self.anchor.device
return MemoryState(
torch.zeros(s["slots"],s["name_bytes"],dtype=torch.uint8,device=dev),
torch.zeros(s["slots"],dtype=torch.long,device=dev),
torch.zeros(s["slots"],s["versions"],s["file_bytes"],dtype=torch.uint8,device=dev),
torch.zeros(s["slots"],s["versions"],dtype=torch.long,device=dev),
torch.zeros(s["slots"],s["versions"],dtype=torch.long,device=dev),
torch.ones(s["slots"],s["versions"],dtype=torch.bool,device=dev),
torch.zeros((),dtype=torch.long,device=dev))
def _name(self,name,device):
b=name.encode("utf-8")
if not b or len(b)>self.spec["name_bytes"] or b"\0" in b:
raise ValueError("Invalid UTF-8 filename length or NUL")
return torch.tensor(list(b.ljust(self.spec["name_bytes"],b"\0")),dtype=torch.uint8,device=device),len(b)
def _slot(self,state,name,create=False):
key,n=self._name(name,state.names.device)
matches=(state.names==key).all(1)&(state.name_length==n)
if bool(matches.any()): return int(matches.long().argmax())
if not create: raise KeyError(name)
free=state.name_length==0
if not bool(free.any()): raise MemoryError("File directory is full")
slot=int(free.long().argmax())
state.names[slot]=key;state.name_length[slot]=n
return slot
def _write(self,state,name,content):
slot=self._slot(state,name,create=content is not None)
payload=b"" if content is None else content.encode("utf-8")
if len(payload)>self.spec["file_bytes"]: raise ValueError("File exceeds byte capacity")
state.clock.add_(1)
sequence=int((state.version[slot]>0).sum())
if sequence<self.spec["versions"]:
column=sequence
else:
column=int(state.version[slot].argmin())
state.data[slot,column].zero_()
if payload:
state.data[slot,column,:len(payload)]=torch.tensor(list(payload),dtype=torch.uint8,device=state.data.device)
state.length[slot,column]=len(payload)
state.version[slot,column]=state.clock
state.deleted[slot,column]=content is None
return int(state.clock)
def forward(self,state,operation,**kwargs):
if operation=="put_many":
files=kwargs["files"]
if not files or len(files)>self.spec["slots"]: raise ValueError("Invalid transaction size")
staged=state.clone()
revisions={}
for name,content in files.items():
if not isinstance(name,str) or (content is not None and not isinstance(content,str)):
raise TypeError("Names and contents must be strings")
revisions[name]=self._write(staged,name,content)
# Commit only after every write passed. No partial externally visible writes.
for name,value in vars(state).items(): value.copy_(getattr(staged,name))
return {"saved":list(files),"revisions":revisions}
if operation=="get":
name=kwargs["name"];slot=self._slot(state,name)
revision=kwargs.get("revision")
if revision is None:
column=int(state.version[slot].argmax())
else:
matches=state.version[slot]==int(revision)
if not bool(matches.any()): raise KeyError("Revision absent or evicted")
column=int(matches.long().argmax())
if bool(state.deleted[slot,column]): raise KeyError(name)
n=int(state.length[slot,column])
return {"name":name,"revision":int(state.version[slot,column]),
"content":bytes(state.data[slot,column,:n].cpu().tolist()).decode("utf-8")}
if operation=="list":
result=[]
for slot in range(self.spec["slots"]):
n=int(state.name_length[slot])
if not n: continue
column=int(state.version[slot].argmax())
if bool(state.deleted[slot,column]): continue
result.append({"name":bytes(state.names[slot,:n].cpu().tolist()).decode("utf-8"),
"revision":int(state.version[slot,column])})
return result
if operation=="rollback":
old=self(state,"get",name=kwargs["name"],revision=kwargs["revision"])
return self(state,"put_many",files={old["name"]:old["content"]})
raise ValueError(f"Unknown memory operation: {operation}")
|