Update app.py

app.py

@@ -1,12 +1,16 @@
 import gradio as gr
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from datasets import load_dataset
 import spaces
 
 # Load model once at startup
-model = GPT2LMHeadModel.from_pretrained("gpt2")
-tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+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
@@ -33,6 +37,9 @@ def train_model():
             loss = outputs.loss
             loss.backward()
             optimizer.step()
+
+        model.push_to_hub("flamiry/first")
+    tokenizer.push_to_hub("flamiry/first")
         
         return f"✅ Training complete! Final Loss: {loss.item():.4f}"
     except Exception as e:
@@ -62,6 +69,4 @@ with gr.Blocks() as demo:
         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.launch()
+        gen_btn.click(generate_text, inputs=prompt_input, outputs=gen_output)