Update app.py

app.py

@@ -73,16 +73,19 @@ def sample_seq(model, context, length, device, temperature=1, top_k=0, top_p=0.0
             next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)
             generated = torch.cat((generated, next_token.unsqueeze(0)), dim=1)
     return generated
-def add_special_tokens():
+def add_special_tokens(lang):
     """ Returns GPT2 tokenizer after adding separator and padding tokens """
-    tokenizer = GPT2Tokenizer.from_pretrained('NlpHUST/gpt2-vietnamese')
+    token = 'gpt2'
+    if lang =='vi':
+        token = 'NlpHUST/gpt2-vietnamese'
+    tokenizer = GPT2Tokenizer.from_pretrained(token)
     special_tokens = {'pad_token': '<|pad|>', 'sep_token': '<|sep|>'}
     tokenizer.add_special_tokens(special_tokens)
     return tokenizer
 
 def gene(t,a):
 
-    tokenizer = add_special_tokens()
+    tokenizer = add_special_tokens(a)
     article = tokenizer.encode(t)[:900]
     # Load model directly
     model = AutoModelForCausalLM.from_pretrained("tiennlu/GPT2en_CNNen_3k")