1	# Models: a step computes with tensors it takes by name from a store, so a
2	# network reads as one step after another and its weights are loaded once.
3	import Base
4	import ./tensor.bend as Tensors
5	import ./tensor_layout.bend as Layout
6	import ./tensor_file.bend as Files
7	
8	def Tensor() -> Type:
9	  Tensors.Tensor<F32>
10	
11	# A store that could not be loaded keeps its problem, and every take reports it.
12	type Store is Type:
13	  Empty{}
14	  Entry{name: String,tensor: Tensor(),rest: Store}
15	  Unavailable{problem: Tensors.Problem}
16	
17	# A step: given the store, a value and the store that remains.
18	def Model(A: Type) -> Type:
19	  Store -> A & Store
20	
21	def Model.pure(-A: Type,value: A) -> Model(A):
22	  store => (value,store)
23	
24	def Model.next(-A: Type,-B: Type,pair: A & Store,next: A -> Model(B)) -> B & Store:
25	  (value,store) = pair
26	  next(value)(store)
27	
28	def Model.bind(-A: Type,-B: Type,step: Model(A),next: A -> Model(B)) -> Model(B):
29	  store => Model.next(A,B,step(store),next)
30	
31	def missing(+name: String) -> Tensor():
32	  Tensors.failed(~F32,"take",Tensors.MissingName{name},[],0.0,1)
33	
34	def unavailable(problem: Tensors.Problem) -> Tensor():
35	  Tensors.Problem{operation,reason} = problem
36	  Tensors.failed(~F32,operation,reason,[],0.0,1)
37	
38	def passed(name: String,tensor: Tensor(),pair: Tensor() & Store) -> Tensor() & Store:
39	  (result,store) = pair
40	  (result,Entry{name,tensor,store})
41	
42	def scan(rest: Store,same: Bool,+name: String,tensor: Tensor(),+wanted: String) -> Tensor() & Store:
43	  match rest same:
44	    case Empty{} True{}: (tensor,Empty{})
45	    case Entry{next,other,more} True{}: (tensor,Entry{next,other,more})
46	    case Unavailable{problem} True{}: (tensor,Unavailable{problem})
47	    case Empty{} False{}: (missing(wanted),Entry{name,tensor,Empty{}})
48	    case Unavailable{+problem} False{}: (unavailable(problem),Entry{name,tensor,Unavailable{problem}})
49	    case Entry{+next,other,more} False{}: passed(name,tensor,scan(more,String.eq(next,wanted),next,other,wanted))
50	
51	def taken(store: Store,+wanted: String) -> Tensor() & Store:
52	  match store:
53	    case Empty{}: (missing(wanted),Empty{})
54	    case Unavailable{+problem}: (unavailable(problem),Unavailable{problem})
55	    case Entry{+name,tensor,rest}: scan(rest,String.eq(name,wanted),name,tensor,wanted)
56	
57	# The tensor of that name leaves the store. A name that is not there gives a
58	# failed tensor.
59	def take(+name: String) -> Model(Tensor()):
60	  store => taken(store,name)
61	
62	def kept(+name: String,pair: Tensor() & Tensor(),store: Store) -> Tensor() & Store:
63	  (tensor,copy) = pair
64	  (tensor,Entry{name,copy,store})
65	
66	# Puts a copy under the name, for a later take, and gives the tensor back.
67	def named_rest(+name: String,tensor: Tensor(),pair: List<&2,String> & Store) -> List<&2,String> & Store:
68	  (names,rest) = pair
69	  (name <> names,Entry{name,tensor,rest})
70	
71	# The names in the store, in its order, and the store unchanged.
72	def names(store: Store) -> List<&2,String> & Store:
73	  match store:
74	    case Empty{}: ([],Empty{})
75	    case Unavailable{problem}: ([],Unavailable{problem})
76	    case Entry{+name,tensor,rest}: named_rest(name,tensor,names(rest))
77	
78	def listed_names() -> Model(List<&2,String>):
79	  store => names(store)
80	
81	def keep(+name: String,tensor: Tensor()) -> Model(Tensor()):
82	  store => kept(name,Layout.copy(~F32,tensor),store)
83	
84	def run(-A: Type,step: Model(A),store: Store) -> A & Store:
85	  step(store)
86	
87	def value(-A: Type,pair: A & Store) -> A:
88	  (result,store) = pair
89	  result
90	
91	# Runs the steps on the store and gives the result; what remains in the store is dropped.
92	def result(-A: Type,step: Model(A),store: Store) -> A:
93	  value(A,step(store))
94	
95	def joined(+name: String,first: Tensor(),second: Tensor(),rest: Store & Store) -> Store & Store:
96	  (left,right) = rest
97	  (Entry{name,first,left},Entry{name,second,right})
98	
99	def copied_entry(+name: String,pair: Tensor() & Tensor(),rest: Store & Store) -> Store & Store:
100	  (first,second) = pair
101	  joined(name,first,second,rest)
102	
103	# Two stores with the same tensors, to run a model more than once.
104	def copy(store: Store) -> Store & Store:
105	  match store:
106	    case Empty{}: (Empty{},Empty{})
107	    case Unavailable{+problem}: (Unavailable{problem},Unavailable{problem})
108	    case Entry{+name,tensor,rest}: copied_entry(name,Layout.copy(~F32,tensor),copy(rest))
109	
110	# Loading: folder/names.txt lists one name per line; folder/<name>.npy holds its tensor.
111	
112	def entry(+name: String,tensor: Tensor(),store: Store) -> Store:
113	  Entry{name,tensor,store}
114	
115	def load_each(names: List<&2,String>,+folder: String) -> IO(Store):
116	  match names:
117	    case Nil{}: IO.pure(Store,Empty{})
118	    case Con{+name,rest}:
119	      match name:
120	        case SNil{}: load_each(rest,folder)
121	        case SCon{head,tail}:
122	          do IO<Store>:
123	            tensor : Tensor() <- Files.load_npy(folder ++ "/" ++ name ++ ".npy")
124	            store : Store <- load_each(rest,folder)
125	            IO.pure(Store,entry(name,tensor,store))
126	
127	def unlisted(+folder: String,detail: String) -> Store:
128	  Unavailable{Tensors.Problem{"load_weights",Tensors.FileFault{folder ++ "/names.txt",detail}}}
129	
130	def named(text: Result<&1,&1,U32 & String,String>,+folder: String) -> IO(Store):
131	  match text:
132	    case Fail{error}: IO.pure(Store,unlisted(folder,"cannot be read"))
133	    case Done{text}: load_each(String.lines(text),folder)
134	
135	def listed(read: File & Result<&1,&1,U32 & String,String>,+folder: String) -> IO(Store):
136	  (file,text) = read
137	  do IO<Store>:
138	    File.close(file)
139	    named(text,folder)
140	
141	def opened(result: Result<&1,&1,U32 & String,File>,+folder: String) -> IO(Store):
142	  match result:
143	    case Fail{error}: IO.pure(Store,unlisted(folder,"cannot be opened"))
144	    case Done{file}:
145	      do IO<Store>:
146	        read : File & Result<&1,&1,U32 & String,String> <- File.read(file,1048576)
147	        listed(read,folder)
148	
149	# A folder whose names.txt cannot be read gives a store whose every take reports it.
150	def load(+folder: String) -> IO(Store):
151	  do IO<Store>:
152	    result : Result<&1,&1,U32 & String,File> <- File.open(folder ++ "/names.txt","r")
153	    opened(result,folder)
154	
155	# Saving: the tensors of a store as folder/<name>.npy, and their names in folder/names.txt.
156	
157	def first(written: File & Result<&1,&1,U32 & String,Unit>) -> File:
158	  (file,result) = written
159	  file
160	
161	def lines(names: List<&2,String>) -> String:
162	  match names:
163	    case []: ""
164	    case name <> rest: name ++ "
165	" ++ lines(rest)
166	
167	def save_each(store: Store,+folder: String) -> IO(Store):
168	  match store:
169	    case Empty{}: IO.pure(Store,Empty{})
170	    case Unavailable{problem}: IO.pure(Store,Unavailable{problem})
171	    case Entry{+name,tensor,rest}:
172	      do IO<Store>:
173	        saved : Tensor() <- Files.save_npy(folder ++ "/" ++ name ++ ".npy",tensor)
174	        others : Store <- save_each(rest,folder)
175	        IO.pure(Store,entry(name,saved,others))
176	
177	def listing(result: Result<&1,&1,U32 & String,File>,text: String) -> IO(Unit):
178	  match result:
179	    case Fail{error}: IO.pure(Unit,Unit{})
180	    case Done{file}:
181	      do IO<Unit>:
182	        written : File & Result<&1,&1,U32 & String,Unit> <- File.write(file,text)
183	        File.close(first(written))
184	
185	def save_listed(pair: List<&2,String> & Store,+folder: String) -> IO(Store):
186	  (names,store) = pair
187	  do IO<Store>:
188	    result : Result<&1,&1,U32 & String,File> <- File.open(folder ++ "/names.txt","w")
189	    listing(result,lines(names))
190	    save_each(store,folder)
191	
192	# Gives the store back; a tensor that could not be written comes back failed.
193	def save(+folder: String,store: Store) -> IO(Store):
194	  save_listed(names(store),folder)