<jupyter_start><jupyter_text>Fit models code<jupyter_code>def AIC(log_likelihood, k):
    """ AIC given log_likelihood and # parameters (k)
    """
    aic = 2 * k - 2 * log_likelihood
    return aic


def BIC(log_likelihood, n, k):
    """ BIC given log_likelihood, number of observations (n) and # parameters (k)
    """
    bic = np.log(n) * k - 2 * log_likelihood
    return bic

def fit_hmm(
    seqs,
    n_components,
    prop_test=0.5,
    verbose=False,
    pom_stop_threshold=1e-8,
    pom_max_iterations=1e8,
):

    # split train and test
    if prop_test == 0:
        seqs_train = seqs_test = seqs
    else:
        # split into train and test for cross validation
        training_mask = np.random.choice(
            np.arange(len(seqs)), size=int(len(seqs) * prop_test), replace=False
        )
        testing_mask = np.array(
            [i for i in np.arange(len(seqs)) if i not in training_mask]
        )

        seqs_train = np.array(seqs)[training_mask]
        seqs_test = np.array(seqs)[testing_mask]

        # make sure test set doesn't contain any data that train doesnt
        assert np.all(
            [
                i in np.unique(np.concatenate(seqs_train))
                for i in np.unique(np.concatenate(seqs_test))
            ]
        )

    # lengths of sequences
    seq_lens = [len(i) for i in seqs_train]

    # get states
    unique_states = np.unique(np.concatenate(seqs_train))

    # model
    hmmlearn_model = hmm.MultinomialHMM(n_components=n_components, verbose=verbose).fit(
        np.concatenate(seqs_train).reshape(-1, 1), [len(i) for i in seqs_train]
    )

    # probability of data given model
    hmmlearn_log_probability = hmmlearn_model.score(
        np.concatenate(seqs_test).reshape(-1, 1), [len(i) for i in seqs_test]
    )

    # hidden labels
    hmmlearn_predictions = [
        hmmlearn_model.predict(np.array(seq).reshape(-1, 1)) for seq in seqs
    ]

    # probability of ending in state, for each state (for pom)
    hmmlearn_end_states = np.array([i[-1] for i in hmmlearn_predictions])
    end_probs = [
        np.sum(hmmlearn_end_states == i)
        / (np.sum(np.concatenate(hmmlearn_predictions) == i) + 1)
        for i in np.arange(n_components)
    ]

    # initialize pom
    transmat = hmmlearn_model.transmat_
    start_probs = hmmlearn_model.startprob_
    dists = hmmlearn_model.emissionprob_

    states = [
        DiscreteDistribution({vis: d[i] for i, vis in enumerate(unique_states)})
        for d in dists
    ]
    pom_model = HiddenMarkovModel.from_matrix(
        transition_probabilities=transmat,
        distributions=states,
        starts=start_probs,
        ends=end_probs,  # discluding ends and merge makes models equal log prob
        merge="None",
    )

    # fit pom model
    improvement = pom_model.fit(
        seqs_train,
        max_iterations=pom_max_iterations,
        verbose=verbose,
        algorithm="baum-welch",
        stop_threshold=pom_stop_threshold,
        n_jobs=1,
    )

    pom_log_probability = np.sum([pom_model.log_probability(seq) for seq in seqs_test])

    pom_predictions = [pom_model.predict(seq) for seq in seqs]

    num_params = pom_model.state_count() + pom_model.edge_count() + pom_model.node_count()
    
    n_data = np.sum([len(i) for i in seqs_test])
    
    aic = AIC(pom_log_probability, num_params)
    bic = BIC(pom_log_probability, n_data, num_params)

    return (
        n_components,
        hmmlearn_model,
        pom_model,
        seqs_train,
        seqs_test,
        hmmlearn_predictions,
        pom_predictions,
        num_params,
        hmmlearn_log_probability,
        pom_log_probability,
        aic,
        bic,
    )
DATASET_ID = 'koumura_bengalese_finch'
embeddings_dfs = list(DATA_DIR.glob('bf_label_dfs/'+DATASET_ID+'/*.pickle'))
DATASET_ID = 'bengalese_finch_sober'
embeddings_dfs = embeddings_dfs + list(DATA_DIR.glob('bf_label_dfs/'+DATASET_ID+'/*.pickle'))
embeddings_dfs[:3]
for loc in tqdm(embeddings_dfs):
    print(indv)
    # read dataframe
    indv_df = pd.read_pickle(loc).sort_values(by=["key", "start_time"])
    indv = indv_df.indv.unique()[0]

    # Get seqs
    hand_seqs = [
        list(indv_df[indv_df.syllables_sequence_id == seqid]["labels_num"].values)
        for seqid in indv_df.syllables_sequence_id.unique()
    ]

    
    n_repeats = 5
    results = Parallel(n_jobs=-1, verbose=15)(
        delayed(fit_hmm)(
            hand_seqs,
            n_components,
            prop_test = 0.0,
            verbose=False,
            pom_stop_threshold=1e-1,
            pom_max_iterations=1e3,
        )
        for n_components in tqdm(
            np.random.permutation(np.repeat(np.arange(15, 45, 2), n_repeats))
        )
    )
    results_df = pd.DataFrame(
        results,
        columns=[
            "n_components",
            "hmmlearn_model",
            "pom_model",
            "seqs_train",
            "seqs_test",
            "hmmlearn_predictions",
            "pom_predictions",
            "n_params",
            "hmmlearn_log_probability",
            "pom_log_probability",
            "aic",
            "bic",
        ],
    )
    
    results_df["indv"] = indv
    save_loc = DATA_DIR / "HMM_fits" / "HMM" / (indv + ".pickle")
    ensure_dir(save_loc)
    results_df.to_pickle(save_loc)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1038475905864774, 'fraction_numerical': 0.006723944714232349, 'mean_word_length': 2.4151785714285716, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '634130', 'license': 'MIT', 'n_tokens_mistral': 1937, 'n_tokens_neox': 1868, 'n_words': 348, 'path': 'notebooks/0.31-compare-sequence-models-bf/0.3-bf-HMM-fit.ipynb', 'repo_name': 'xingjeffrey/avgn_paper', 'size': '5354'}