Update app.py

app.py

@@ -11,16 +11,13 @@ login(token=os.environ.get("hf_token"))
 try:
     model = AutoModelForCausalLM.from_pretrained("flamiry/first")
     tokenizer = AutoTokenizer.from_pretrained("flamiry/first")
-except:
-    model = AutoModelForCausalLM.from_pretrained("gpt2")
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
 
 @spaces.GPU
-def train_model():
+def train_model(start, end):
     try:
         dataset = load_dataset("DGurgurov/slovak_sa", split="train")
-        slovak_texts = list(dataset['text'])[:50]
+        slovak_texts = list(dataset['text'])[start:end]
         
         inputs = tokenizer(
             slovak_texts,
@@ -55,17 +52,19 @@ def generate_text(prompt):
         return f"❌ Error: {str(e)}"
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🇸🇰 Slovak LLM Training")
+    gr.Markdown("#Slovak LLM Training")
     
     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")
         train_btn = gr.Button("Start Training", variant="primary")
         train_output = gr.Textbox(label="Result", interactive=False)
-        train_btn.click(train_model, 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")
-        prompt_input = gr.Textbox(label="Prompt", placeholder="Mačka je...")
+        prompt_input = gr.Textbox(label="Prompt")
         gen_btn = gr.Button("Generate")
         gen_output = gr.Textbox(label="Generated Text", interactive=False)
         gen_btn.click(generate_text, inputs=prompt_input, outputs=gen_output)