INNER CODE UNIT · Python

clean_keywords

gagan3012/keytotext · keytotext/dataset.py:12

def clean_keywords(keywords):
  return clean(list(map(list, zip(*keywords)))[0])

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"]

View source record →