Update app.py

app.py

@@ -13,12 +13,11 @@ tokenizer = AutoTokenizer.from_pretrained("flamiry/first")
 tokenizer.pad_token = tokenizer.eos_token
 
 @spaces.GPU
-def train_model(start, end):
-    start = int(start)
-    end = int(end)
+def train_model(number):
+    number = int(number)
     try:
-        dataset = load_dataset("DGurgurov/slovak_sa", split="train")
-        slovak_texts = list(dataset['text'])[start:end]
+        dataset = load_dataset("allenai/c4", "sk", split="train", streaming=True)
+        slovak_texts = list(dataset.take(number))
         
         inputs = tokenizer(
             slovak_texts,
@@ -57,8 +56,7 @@ with gr.Blocks() as demo:
     
     with gr.Tab("Train Model"):
         gr.Markdown("Click to train the model on Slovak data")
-        start_input = gr.Textbox(label="Start")
-        end_input = gr.Textbox(label="End")
+        number_input = gr.Textbox(label="Kol-vo textov")
         train_btn = gr.Button("Start Training", variant="primary")
         train_output = gr.Textbox(label="Result", interactive=False)
         train_btn.click(train_model, inputs=[start_input,end_input] ,outputs=train_output)