vivianhuang88/bert_twitter_hashtag

Overview

This model is based on bert-base-uncased model and trained on more than 30k tweets that scraped from Twitter. By inputing some sentences with a '[MASK]' indicating the location you would like to fill in with a hashtag, our model can generate potential related trending topics according to your tweet context.

Define a list of trending topics

trending_topics = [Your choice of topics]

Download the model

from transformers import pipeline, BertTokenizer 
import numpy as np

MODEL = "vivianhuang88/bert_twitter_hashtag"
fill_mask = pipeline("fill-mask", model=MODEL, tokenizer=MODEL)
tokenizer = BertTokenizer.from_pretrained(MODEL, additional_special_tokens=trending_topics)

Get the output

def print_candidates(text, candidates):
    for i in range(5):
        token = tokenizer.decode(candidates[i]['token'])
        topic = ''.join(token.split())
        output = text.replace("[MASK]", topic)
        print(output)
        
text = "Bruce has an electric guitar set in [MASK]. "
candidates = fill_mask(text, targets = trending_topics)
print_candidates(text, candidates)