Update app.py

app.py

@@ -1,7 +1,67 @@
 import gradio as gr
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+import torch
+from datasets import load_dataset
+import spaces
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Load model once at startup
+model = GPT2LMHeadModel.from_pretrained("gpt2")
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token
+
+@spaces.GPU
+def train_model():
+    try:
+        # Load Slovak data
+        dataset = load_dataset("DGurgurov/slovak_sa", split="train")
+        slovak_texts = list(dataset['text'])[:200]  # Only 200 texts
+        
+        # Tokenize
+        inputs = tokenizer(
+            slovak_texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=128
+        )
+        
+        # Train
+        optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
+        for epoch in range(2):
+            optimizer.zero_grad()
+            outputs = model(**inputs, labels=inputs['input_ids'])
+            loss = outputs.loss
+            loss.backward()
+            optimizer.step()
+        
+        return f"✅ Training complete! Final Loss: {loss.item():.4f}"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+@spaces.GPU
+def generate_text(prompt):
+    try:
+        input_ids = tokenizer.encode(prompt, return_tensors="pt")
+        output = model.generate(input_ids, max_length=50)
+        return tokenizer.decode(output[0], skip_special_tokens=True)
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+# Create interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 🇸🇰 Slovak LLM Training")
+    
+    with gr.Tab("Train Model"):
+        gr.Markdown("Click to train the model on Slovak data")
+        train_btn = gr.Button("Start Training", variant="primary")
+        train_output = gr.Textbox(label="Result", interactive=False)
+        train_btn.click(train_model, outputs=train_output)
+    
+    with gr.Tab("Generate Text"):
+        gr.Markdown("Generate Slovak text")
+        prompt_input = gr.Textbox(label="Prompt", placeholder="Mačka je...")
+        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)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()