class WrappedRiverClusterAlgo:
    """Wraps a River model so that it can be used to train the model in chunks of data similar
    to online training
    """
    def __init__(self, model):
        self.model = model

    def partial_fit(self, umap_embeddings):
        for umap_embedding, _ in stream.iter_array(umap_embeddings):
            self.model = self.model.learn_one(umap_embedding)

        labels = []
        for umap_embedding, _ in stream.iter_array(umap_embeddings):
            label = self.model.predict_one(umap_embedding)
            labels.append(label)
    
        self.labels_ = labels
        return self
        
# Step 1 - Extract embeddings           
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')                       

# Step 2 - Reduce dimensionality       
umap_model = IncrementalPCA(n_components=5)

# Step 3 - Cluster reduced embeddings             
cluster_model = WrappedRiverClusterAlgo(cluster.CluStream())

# Step 4 - Tokenize topics          
vectorizer_model = OnlineCountVectorizer(stop_words="english", decay=.01, delete_min_df=10.00,
                                                 ngram_range=(2,2))        

# Step 5 - Create topic representation
ctfidf_model = ClassTfidfTransformer(reduce_frequent_words=True)

representation_model = KeyBERTInspired(nr_repr_docs=15,random_state=100)

# All steps together
topic_model = BERTopic(
        embedding_model=embedding_model,          
        umap_model=umap_model,                   
        hdbscan_model=cluster_model,              
        vectorizer_model=vectorizer_model,       
        ctfidf_model=ctfidf_model,
        calculate_probabilities=True,
        representation_model=representation_model,
        nr_topics="auto",
        verbose=True)


for data in dataset:    
    topic_model.partial_fit(data)  
    topics.extend(topic_model.topics_)      
        
# Update model topics attribute
topic_model.topics_ = topics
        
# Save the model            
topic_model.save(model_safatensors_path,  serialization="safetensors", save_ctfidf=True,
                         save_embedding_model="sentence-transformers/all-MiniLM-L6-v2")