Error when specifying index_clustering_model in predict

[tseidl]

Hi guys,

First off, I'd like to commend you on the cool paper and package. I've encountered an issue concerning the predict method, specifically regarding the ability to manually specify the index_clustering_model.

When I don't specify index_clustering_model, everything runs smoothly:

narratives = m.predict(postproc_roles, index_clustering_model=None, progress_bar=True)

However, when I try specifying it, for instance:

narratives = m.predict(postproc_roles, index_clustering_model=3, progress_bar=True)

I encounter the following error:

UnboundLocalError: local variable 'clustering_model' referenced before assignment

Just from quickly going over the the code, it looks like the issue stems from the conditional assignment of the clustering_model variable. Specifically:

If index_clustering_model is set to None, then clustering_model is never initialized, which results in the mentioned error when the function later tries to access clustering_model.

Not sure if I'm overlooking something, but if that's the case, one would just need to modify the function as follows:

def predict(self, srl_res, index_clustering_model=None, progress_bar=False):
    # ... other parts of the function ...

    # Ensure clustering_model is assigned correctly:
    if index_clustering_model is None:
        clustering_model = self.clustering_models[self.index_optimal_model]
    else:
        clustering_model = self.clustering_models[index_clustering_model]

    # ... rest of the function ...

I'm somewhat new to python, so this might be off, but thought I'd ask/let you know. Keep up the good work!