mirror of
https://github.com/Hestia-Homes/Model.git
synced 2026-06-08 11:17:27 +00:00
21 lines
No EOL
580 B
Python
21 lines
No EOL
580 B
Python
import pandas as pd
|
|
from core.Logger import logger
|
|
|
|
class DataLoader():
|
|
|
|
@staticmethod
|
|
def load(filepath: str, index_col: str = None) -> pd.DataFrame:
|
|
"""
|
|
Load different datasets
|
|
"""
|
|
if filepath.endswith('.parquet'):
|
|
df = pd.read_parquet(filepath)
|
|
if index_col is not None:
|
|
df = df.set_index(index_col)
|
|
elif filepath.endswith('.csv'):
|
|
df = pd.read_csv(filepath, index_col=index_col)
|
|
else:
|
|
logger.error('Not implemented!')
|
|
exit(1)
|
|
|
|
return df |