Account for un-approx variances

reglab · Nov 10, 2024 · 0ceb02b · 0ceb02b
1 parent 2667010
commit 0ceb02b
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/glove_v/variance.py b/glove_v/variance.py
@@ -40,7 +40,7 @@ def load_variance(
                 diagonal = f.get_tensor(f"diag_{word_idx}")
                 return np.diag(diagonal)
 
-            # Otherwise, it must be an SVD approximation
+            # SVD approximation
             elif f"U_{word_idx}" in f.keys():
                 U = f.get_tensor(f"U_{word_idx}")
                 s = f.get_tensor(f"s_{word_idx}")
@@ -49,6 +49,11 @@ def load_variance(
                 # Reconstruct using SVD components: U * diag(s) * Vt
                 return U @ np.diag(s) @ Vt
 
+            # Complete approximation (SVD unavailable)
+            elif f"complete_{word_idx}" in f.keys():
+                complete = f.get_tensor(f"complete_{word_idx}")
+                return complete
+
             else:
                 raise KeyError(
                     f"[ERROR No approximation found for word index {word_idx}"