utils/ml_model.py

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier


def train_model(counts: pd.DataFrame, labels: list):

    """
    Train Random Forest model on RNA-seq data

    counts:
        rows = samples
        columns = genes

    labels:
        sample classes
    """

    X = counts
    y = np.array(labels)


    model = RandomForestClassifier(
        n_estimators=100,
        random_state=42
    )

    model.fit(X, y)


    feature_importances = pd.DataFrame({

        "Gene": counts.columns,

        "Importance": model.feature_importances_

    }).sort_values(
        by="Importance",
        ascending=False
    )


    return model, feature_importances



def get_top_biomarkers(
        feature_importances: pd.DataFrame,
        top_n=10
):

    return feature_importances.head(top_n)