Update app.py

app.py

@@ -3,13 +3,19 @@ import gradio as gr
 
 import torch
 
+from datasets import load_dataset
+
+
+
 first_generation = True
 prefix = ''
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 model_checkpoint = "hackathon-pln-es/es_text_neutralizer"
+data_checkpoint = "hackathon-pln-es/neutral-es"
 
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+dataset = load_dataset(data_checkpoint, split='test')
 
 model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
 # TODO: jarl!! check this for avoiding short segments