Adding pre-trained bert

app.py

@@ -2,14 +2,22 @@ import gradio as gr
 from NeuralTextGenerator import BertTextGenerator
 
 # Load models
-model_name = "cardiffnlp/twitter-xlm-roberta-base"
-en_model = BertTextGenerator(model_name, tokenizer=model_name)
 
+## BERT
+BERT_model_name = "Twitter/twhin-bert-large"
+BERT = BertTextGenerator(BERT_model_name, tokenizer=BERT_model_name)
+
+## RoBERTa
+RoBERTa_model_name = "cardiffnlp/twitter-xlm-roberta-base"
+RoBERTa = BertTextGenerator(RoBERTa_model_name, tokenizer=RoBERTa_model_name)
+
+## Finetuned BERT
 finetunned_BERT_model_name = "JuanJoseMV/BERT_text_gen"
-finetunned_BERT_en_model = BertTextGenerator(finetunned_BERT_model_name, tokenizer='bert-base-uncased')
+finetunned_BERT = BertTextGenerator(finetunned_BERT_model_name, tokenizer='bert-base-uncased')
 
+## Finetuned RoBERTa
 finetunned_RoBERTa_model_name = "JuanJoseMV/XLM_RoBERTa_text_gen"
-finetunned_RoBERTa_en_model = BertTextGenerator(finetunned_RoBERTa_model_name, tokenizer=finetunned_RoBERTa_model_name)
+finetunned_RoBERTa = BertTextGenerator(finetunned_RoBERTa_model_name, tokenizer=finetunned_RoBERTa_model_name)
 
 special_tokens = [
     '[POSITIVE-0]', 
@@ -20,23 +28,19 @@ special_tokens = [
     '[NEGATIVE-2]'
     ]
 
-# en_model.tokenizer.add_special_tokens({'additional_special_tokens': special_tokens})
-# en_model.model.resize_token_embeddings(len(en_model.tokenizer))
-
-finetunned_BERT_en_model.tokenizer.add_special_tokens({'additional_special_tokens': special_tokens})
-finetunned_BERT_en_model.model.resize_token_embeddings(len(en_model.tokenizer))
-
-# finetunned_RoBERTa_en_model.tokenizer.add_special_tokens({'additional_special_tokens': special_tokens})
-# finetunned_RoBERTa_en_model.model.resize_token_embeddings(len(en_model.tokenizer))
+finetunned_BERT.tokenizer.add_special_tokens({'additional_special_tokens': special_tokens})
+finetunned_BERT.model.resize_token_embeddings(len(finetunned_BERT.tokenizer))
 
 def sentence_builder(selected_model, n_sentences, max_iter, sentiment, seed_text):
 
-    if selected_model == "Finetuned_RoBERTA":
-        generator = finetunned_RoBERTa_en_model
+    if selected_model == "Finetuned_RoBERTa":
+        generator = finetunned_RoBERTa
     elif selected_model == "Finetuned_BERT":
-        generator = finetunned_BERT_en_model
+        generator = finetunned_BERT
+    elif selected_model == "RoBERTa":
+        generator = RoBERTa
     else:
-        generator = en_model
+        generator = BERT
 
     parameters = {'n_sentences': n_sentences,  
               'batch_size': 2,
@@ -63,7 +67,7 @@ def sentence_builder(selected_model, n_sentences, max_iter, sentiment, seed_text
 demo = gr.Interface(
     sentence_builder,
     [
-        gr.Radio(["Pre-trained", "Finetuned_RoBERTA", "Finetunned_BERT"], value="Pre-trained", label="Sentiment to generate"),
+        gr.Radio(["BERT", "RoBERTa", "Finetuned_RoBERTa", "Finetunned_BERT"], value="BERT", label="Generator model"),
         gr.Slider(1, 15, value=2, label="Num. Tweets", step=1, info="Number of tweets to be generated."),
         gr.Slider(50, 500, value=100, label="Max. iter", info="Maximum number of iterations for the generation."),
         gr.Radio(["POSITIVE", "NEGATIVE"], value="POSITIVE", label="Sentiment to generate"),