fix some frequency handling issues

ibm-granite · Jul 31, 2024 · 5fb1cd0 · 5fb1cd0
1 parent 9b4801d
commit 5fb1cd0
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tsfm_public/toolkit/util.py b/tsfm_public/toolkit/util.py
@@ -565,16 +565,19 @@ def convert_tsf(filename: str) -> pd.DataFrame:
         "seconds": "s",
         "minutes": "min",
         "minutely": "min",
+        "monthly": "MS",
+        "yearly": "YS",
+        "quarterly": "QS",
     }
 
     if frequency:
         try:
             freq_val, freq_unit = frequency.split("_")
             freq = freq_val + tsf_to_pandas_freq_map[freq_unit]
         except ValueError:
-            freq = tsf_to_pandas_freq_map(freq_val)
+            freq = tsf_to_pandas_freq_map[frequency]
         except KeyError:
-            raise ValueError(f"Input file contains an unknow frequency unit {freq_unit}")
+            raise ValueError(f"Input file contains an unknown frequency unit {freq_unit}")
     else:
         freq = None
 
@@ -596,6 +599,7 @@ def convert_tsf(filename: str) -> pd.DataFrame:
         )
 
     df = pd.concat(dfs)
+    df.reset_index(inplace=True, drop=True)
     return df