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


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):
    model = load_model(particle)

    # Static input values
    input_data = pd.DataFrame(
        {"Temperature": [20.0], "Wind Speed": [10.0], "Humidity": [50.0]}
    )

    # Run the model with static input
    prediction = model.predict(input_data)
    return prediction