Start-GPT commited on
Commit
9248a10
1 Parent(s): efac4b5

Create gen_utils.py

Browse files
Files changed (1) hide show
  1. server/utils/gen_utils.py +54 -0
server/utils/gen_utils.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from copy import deepcopy
2
+ import numpy as np
3
+ from functools import partial
4
+ from .f import memoize
5
+
6
+ def check_key_len(d, length):
7
+ for k, v in d.items():
8
+ if len(v) != length:
9
+ raise ValueError(f"dictionary values are not all of length {length}. Found {len(v)}")
10
+
11
+ def check_zippable(dict_a, dict_b):
12
+ """Check that the arrays contained in each value of dict_a and b are of identical length"""
13
+ avals = list(dict_a.values())
14
+ bvals = list(dict_b.values())
15
+
16
+ assert len(avals) > 0
17
+
18
+ length = len(avals[0])
19
+ check_key_len(dict_a, length)
20
+ check_key_len(dict_b, length)
21
+
22
+ def zip_dicts(dict_a, dict_b):
23
+ """Zip the arrays associated with the keys in two dictionaries"""
24
+ combined = {}
25
+ combined.update(dict_a)
26
+ combined.update(dict_b)
27
+ zipped_vals = zip(*combined.values())
28
+ keys = list(combined.keys())
29
+
30
+ out = []
31
+ for i, zv in enumerate(zipped_vals):
32
+ obj = {k: v_ for (k,v_) in zip(keys, zv)}
33
+ out.append(obj)
34
+
35
+ return out
36
+
37
+ def vround(ndigits):
38
+ """Vectorized version of "round" that can be used on numpy arrays. Returns a function that can be used to round digits in a response"""
39
+ return np.vectorize(partial(round, ndigits=ndigits))
40
+
41
+ def roundTo(arr, ndigits):
42
+ """Round an array to ndigits"""
43
+ f = vround(ndigits)
44
+ return f(arr)
45
+
46
+ def map_nlist(f, nlist):
47
+ """Map a function across an arbitrarily nested list"""
48
+ new_list=[]
49
+ for i in range(len(nlist)):
50
+ if isinstance(nlist[i],list):
51
+ new_list += [map_nlist(f, nlist[i])]
52
+ else:
53
+ new_list += [f(nlist[i])]
54
+ return new_list