将主题建模结果投射到数据框

Question

. Answers to this question are eligible for a +150 reputation bounty. xavi is looking for a .

I have used BertTopic with KeyBERT to extract some topics from some docs

from bertopic import BERTopic
topic_model = BERTopic(nr_topics="auto", verbose=True, n_gram_range=(1, 4), calculate_probabilities=True, embedding_model='paraphrase-MiniLM-L3-v2', min_topic_size= 3)
topics, probs = topic_model.fit_transform(docs)

Now I can access the topic name

freq = topic_model.get_topic_info()
print("Number of topics: {}".format( len(freq)))
freq.head(30)

   Topic    Count   Name
0   -1       1     -1_default_greenbone_gmp_manager
1    0      14      0_http_tls_ssl tls_ssl
2    1      8       1_jboss_console_web_application

and inspect the topics

[('http', 0.0855701486234524),          
 ('tls', 0.061977919455444744),
 ('ssl tls', 0.061977919455444744),
 ('ssl', 0.061977919455444744),
 ('tcp', 0.04551718585531556),
 ('number', 0.04551718585531556)]

[('jboss', 0.14014705432060262),
 ('console', 0.09285308122803233),
 ('web', 0.07323749337563096),
 ('application', 0.0622930523123512),
 ('management', 0.0622930523123512),
 ('apache', 0.05032395169459188)]

What I want is to have a final data frame that has in one column the topic name and in another column the elements of the topic

expected outcome:

  class                         entities
o http_tls_ssl tls_ssl           HTTP...etc
1 jboss_console_web_application  JBoss, console, etc

and one dataframe with the topic name on different columns

  http_tls_ssl tls_ssl           jboss_console_web_application
o http                           JBoss
1 tls                            console
2 etc                            etc

I did not find out how to do this. Is there a way?

Answer 1

Here is one way to to it:

Setup

import pandas as pd
from bertopic import BERTopic
from sklearn.datasets import fetch_20newsgroups

docs = fetch_20newsgroups(subset="all", remove=("headers", "footers", "quotes"))["data"]

topic_model = BERTopic()
topics, probs = topic_model.fit_transform(docs[:1_000])

df = topic_model.get_topic_info()
print(df)
# Output
   Topic  Count                          Name
0      0    875               0_the_to_of_and
1      1     93               1_the_to_and_in
2      2     32  2_testing_deletion_hello_was

First dataframe

Using Pandas string methods :

df = df.rename(columns={"Name": "class"}).drop(columns=["Topic", "Count"])
df["class"] = df["class"].str.replace("-", "").apply(lambda x: x[2:])  # remove '-1_',...
df["entities"] = df["class"].str.split("_")

print(df)
# Output
                        class                         entities
0               the_to_of_and               [the, to, of, and]
1               the_to_and_in               [the, to, and, in]
2  testing_deletion_hello_was  [testing, deletion, hello, was]

Second dataframe

Using Pandas transpose :

other_df = df.T.reset_index(drop=True)
new_col_labels = other_df.iloc[0]  # save first row
other_df = other_df[1:]  # remove first row
other_df.columns = new_col_labels
other_df = pd.DataFrame({col: other_df.loc[1, col] for col in other_df.columns})

print(other_df)
# Output
  the_to_of_and the_to_and_in testing_deletion_hello_was
0           the           the                    testing
1            to            to                   deletion
2            of           and                      hello
3           and            in                        was