Update app.py

app.py

@@ -9,10 +9,15 @@ from itertools import islice
 
 login(token=os.environ.get("hf_token"))
 
+# Load model and tokenizer
 model = AutoModelForCausalLM.from_pretrained("flamiry/first")
 tokenizer = AutoTokenizer.from_pretrained("flamiry/first")
 tokenizer.pad_token = tokenizer.eos_token
 
+# Move model to GPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
 @spaces.GPU
 def train_model(start, end):
     start = int(start)
@@ -28,6 +33,9 @@ def train_model(start, end):
             max_length=512
         )
         
+        # Move inputs to device
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        
         optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
         for epoch in range(2):
             optimizer.zero_grad()
@@ -35,7 +43,7 @@ def train_model(start, end):
             loss = outputs.loss
             loss.backward()
             optimizer.step()
-
+        
         model.push_to_hub("flamiry/first")
         tokenizer.push_to_hub("flamiry/first")
         
@@ -46,14 +54,14 @@ def train_model(start, end):
 @spaces.GPU
 def generate_text(prompt):
     try:
-        input_ids = tokenizer.encode(prompt, return_tensors="pt")
+        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
         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)}"
 
 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")
@@ -61,7 +69,7 @@ with gr.Blocks() as demo:
         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=[start_input, end_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")