INNER CODE UNIT · Python
make_keywords
gagan3012/keytotext · keytotext/dataset.py:15
def make_keywords(dataset):
kw_model = KeyBERT()
df = pd.DataFrame(columns=['text', 'keywords'])
df["text"] = dataset["text"]
for i in tqdm(range(len(df))):
keyword = kw_model.extract_keywords(df['text'][i])
clean = clean_keywords(keyword)
df["keywords"][i] = clean
return df
def make_dataset(dataset="common_gen", split="train"):
if dataset == "common_gen":
dataset = load_dataset(dataset, split=split)
df = pd.DataFrame()
df["keywords"] = dataset["concepts"]
df["text"] = dataset["target"]
df["keywords"] = df["keywords"].apply(clean)
return df