src/models_loading.py

import os

import joblib
import pandas as pd
import streamlit as st
from dotenv import load_dotenv
from huggingface_hub import hf_hub_download, login
from src.data_loading import create_features

def load_model(particle):
    load_dotenv()
    login(token=os.getenv("HUGGINGFACE_DOWNLOAD_TOKEN"))

    repo_id = f"elisaklunder/Utrecht-{particle}-Forecasting-Model"
    if particle == "O3":
        file_name = "O3_svr_model.pkl"
    elif particle == "NO2":
        file_name == "hehehe"

    model_path = hf_hub_download(repo_id=repo_id, filename=file_name)
    model = joblib.load(model_path)

    return model


@st.cache_resource(ttl=6 * 300)  # Reruns every 6 hours
def run_model(particle, data):
    input_data = create_features(data=data, target_particle=particle)
    model = load_model(particle)

    # Run the model with static input
    prediction = model.predict(input_data)
    target_scaler = joblib.load(f"scalers/target_scaler_{particle}.joblib")
    prediction = target_scaler.inverse_transform(prediction)
    return prediction