Update app.py

app.py

@@ -1,20 +1,21 @@
-import streamlit as st
+import os
 import json
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
+import streamlit as st
 from datasets import Dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
 from peft import LoraConfig, get_peft_model
-from huggingface_hub import HfApi, HfFolder, Repository
+from huggingface_hub import Repository
 
 # -------- CONFIG ----------
 MODEL_ID = "Neon-AI/Niche"
 CHECKPOINT_DIR = "./checkpoints"
-HF_TOKEN = st.secrets["HF_TOKEN"]
+HF_TOKEN = st.secrets["HF_TOKEN"]  # Put your HF token in Streamlit secrets
 
 st.title("🧠 Niche Trainer with Push to HF")
 
 # ---------- Load model once ----------
-@st.cache_resource
+@st.cache_resource(show_spinner=True)
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
     model = AutoModelForCausalLM.from_pretrained(
@@ -39,6 +40,9 @@ json_input = st.text_area(
     placeholder='[{"prompt": "...", "response": "..."}]'
 )
 
+# ---------- Max token length ----------
+max_len = st.slider("Max token length", min_value=64, max_value=512, value=256)
+
 # ---------- Train ----------
 train_started = False
 if st.button("Train"):
@@ -54,7 +58,7 @@ if st.button("Train"):
             ds = Dataset.from_dict({"text": texts})
 
             def tokenize(batch):
-                out = tokenizer(batch["text"], truncation=True, padding="max_length", max_length=256)
+                out = tokenizer(batch["text"], truncation=True, padding="max_length", max_length=max_len)
                 out["labels"] = out["input_ids"].copy()
                 return out
 
@@ -70,8 +74,7 @@ if st.button("Train"):
                     lora_dropout=0.1,
                     target_modules=["c_attn"]
                 )
-                model_peft = get_peft_model(model, peft_config)
-                train_model = model_peft
+                train_model = get_peft_model(model, peft_config)
             else:
                 train_model = model
 
@@ -96,23 +99,30 @@ if st.button("Train"):
             trainer.train()
             st.success("✅ Training done!")
             train_started = True
+
+            # Use trained model for chat
+            model = train_model
+
     except Exception as e:
-        st.error(f"Error: {e}")
+        st.error(f"Error during training: {e}")
 
 # ---------- Push to HF ----------
 if train_started and st.button("Push to Hugging Face"):
     try:
-        repo = Repository(
-            local_dir=CHECKPOINT_DIR,
-            clone_from=MODEL_ID,
-            use_auth_token=HF_TOKEN
-        )
+        # Prepare repo
+        if os.path.exists(CHECKPOINT_DIR):
+            repo = Repository(local_dir=CHECKPOINT_DIR, use_auth_token=HF_TOKEN)
+        else:
+            repo = Repository(local_dir=CHECKPOINT_DIR, clone_from=MODEL_ID, use_auth_token=HF_TOKEN)
+
         # Save trained model + tokenizer
-        train_model.save_pretrained(CHECKPOINT_DIR)
+        model.save_pretrained(CHECKPOINT_DIR)
         tokenizer.save_pretrained(CHECKPOINT_DIR)
 
+        # Push
         repo.push_to_hub(commit_message="Update Niche model with new training")
         st.success("✅ Model pushed to HF successfully!")
+
     except Exception as e:
         st.error(f"Push failed: {e}")