Update app.py

app.py

@@ -5,6 +5,7 @@ from datasets import load_dataset
 import spaces
 from huggingface_hub import login
 import os
+from itertools import islice
 
 login(token=os.environ.get("hf_token"))
 
@@ -13,17 +14,18 @@ tokenizer = AutoTokenizer.from_pretrained("flamiry/first")
 tokenizer.pad_token = tokenizer.eos_token
 
 @spaces.GPU
-def train_model(number):
-    number = int(number)
+def train_model(start, end):
+    start = int(start)
+    end = int(end)
     try:
         dataset = load_dataset("allenai/c4", "sk", split="train", streaming=True)
-        slovak_texts = [example['text'] for example in dataset.take(number)]        
+        slovak_texts = [example['text'] for example in islice(dataset, start, end)]
         inputs = tokenizer(
             slovak_texts,
             return_tensors="pt",
             padding=True,
             truncation=True,
-            max_length=64
+            max_length=512
         )
         
         optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
@@ -55,10 +57,11 @@ with gr.Blocks() as demo:
     
     with gr.Tab("Train Model"):
         gr.Markdown("Click to train the model on Slovak data")
-        number_input = gr.Textbox(label="Kol-vo textov")
+        start_input = gr.Textbox(label="Start")
+        end_input = gr.Textbox(label="End")
         train_btn = gr.Button("Start Training", variant="primary")
         train_output = gr.Textbox(label="Result", interactive=False)
-        train_btn.click(train_model, inputs=[number_input] ,outputs=train_output)
+        train_btn.click(train_model, inputs=[start_input, end_input] ,outputs=train_output)
     
     with gr.Tab("Generate Text"):
         gr.Markdown("Generate Slovak text")