prepforge 0.2.1__tar.gz → 0.2.2__tar.gz

{prepforge-0.2.1 → prepforge-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prepforge
-Version: 0.2.1
+Version: 0.2.2
 Summary: AI Exam Preparation System with TUI, GUI, and Streamlit
 Author-email: Himanshu Arora <ha0himanshuarora@gmail.com>
 License-Expression: MIT

prepforge-0.2.2/major_project/core/train.py

@@ -0,0 +1,156 @@
+from unsloth import FastLanguageModel
+from datasets import load_dataset, Dataset
+from transformers import TrainingArguments
+from trl import SFTTrainer
+
+from major_project.config import get_model
+
+
+def train_model(
+    dataset_path,
+    output_dir,
+    epochs=2,
+    lr=2e-4,
+    batch_size=2,
+    max_seq_length=512,
+    limit=None,
+    subset=None,
+):
+    # -----------------------------
+    # LOAD MODEL FROM CONFIG
+    # -----------------------------
+    model_name = get_model()
+    print(f"[TRAIN] Using model: {model_name}")
+
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=model_name,
+        max_seq_length=max_seq_length,
+        load_in_4bit=True,
+    )
+
+    # -----------------------------
+    # APPLY LoRA
+    # -----------------------------
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=16,
+        target_modules=[
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        lora_alpha=16,
+        lora_dropout=0,
+        bias="none",
+        use_gradient_checkpointing="unsloth",
+    )
+
+    # -----------------------------
+    # LOAD DATASET
+    # -----------------------------
+    dataset = load_dataset("json", data_files=dataset_path, split="train")
+    assert isinstance(dataset, Dataset)
+
+    print(f"[TRAIN] Original dataset size: {len(dataset)}")
+
+    # -----------------------------
+    # APPLY LIMIT / SUBSET
+    # -----------------------------
+    if limit is not None:
+        dataset = dataset.select(range(min(limit, len(dataset))))
+        print(f"[TRAIN] Using first {len(dataset)} samples")
+
+    elif subset is not None:
+        size = int(len(dataset) * subset)
+        dataset = dataset.select(range(size))
+        print(f"[TRAIN] Using {subset * 100:.1f}% → {len(dataset)} samples")
+
+    # -----------------------------
+    # FORMAT DATA
+    # -----------------------------
+    def format_prompt(example):
+        instruction = example.get("instruction", "")
+        input_text = example.get("input", "")
+        output = example.get("output", "")
+
+        prompt = f"{instruction}\n{input_text}".strip()
+
+        user_part = f"""<|begin_of_text|>
+<|start_header_id|>user<|end_header_id|>
+
+{prompt}
+
+<|start_header_id|>assistant<|end_header_id|>
+
+"""
+
+        full_text = user_part + output
+
+        user_tokens = tokenizer(user_part, add_special_tokens=False)
+        full_tokens = tokenizer(
+            full_text,
+            truncation=True,
+            max_length=max_seq_length,
+        )
+
+        labels = full_tokens["input_ids"].copy()
+        user_len = len(user_tokens["input_ids"])
+
+        # Mask user part
+        labels[:user_len] = [-100] * user_len
+
+        return {
+            "input_ids": full_tokens["input_ids"],
+            "attention_mask": full_tokens["attention_mask"],
+            "labels": labels,
+        }
+
+    dataset = dataset.map(format_prompt, batched=False)
+
+    # Keep only required columns
+    dataset = dataset.remove_columns(
+        [
+            col
+            for col in dataset.column_names
+            if col not in ["input_ids", "attention_mask", "labels"]
+        ]
+    )
+
+    # -----------------------------
+    # TRAINER
+    # -----------------------------
+    training_args = TrainingArguments(
+        per_device_train_batch_size=batch_size,
+        num_train_epochs=epochs,
+        learning_rate=lr,
+        output_dir=output_dir,
+        logging_steps=10,
+        save_steps=500,
+        optim="adamw_8bit",
+        report_to="none",  # avoids wandb issues
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=dataset,
+        args=training_args,
+    )
+
+    # -----------------------------
+    # TRAIN
+    # -----------------------------
+    print("[TRAIN] Starting training...")
+    trainer.train()
+
+    # -----------------------------
+    # SAVE
+    # -----------------------------
+    model.save_pretrained(output_dir)
+    tokenizer.save_pretrained(output_dir)
+
+    print("[TRAIN] Training complete.")

{prepforge-0.2.1 → prepforge-0.2.2}/prepforge.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prepforge
-Version: 0.2.1
+Version: 0.2.2
 Summary: AI Exam Preparation System with TUI, GUI, and Streamlit
 Author-email: Himanshu Arora <ha0himanshuarora@gmail.com>
 License-Expression: MIT

{prepforge-0.2.1 → prepforge-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "prepforge"
-version = "0.2.1"
+version = "0.2.2"
 description = "AI Exam Preparation System with TUI, GUI, and Streamlit"
 authors = [
     { name = "Himanshu Arora", email = "ha0himanshuarora@gmail.com" }

prepforge-0.2.1/major_project/core/train.py

@@ -1,212 +0,0 @@
-import os
-
-# -----------------------------
-# 🔥 FORCE OFFLINE MODE
-# -----------------------------
-os.environ["TRANSFORMERS_OFFLINE"] = "1"
-os.environ["HF_HUB_OFFLINE"] = "1"
-
-import streamlit as st
-
-# 🔥 CORE IMPORTS
-from major_project.core.model import load_model
-from major_project.core.utils import (
-    build_prompt,
-    generate_stream,
-    chat_history,
-    last_mcq_block,
-)
-from major_project.config import DEFAULT_APP_NAME, DEFAULT_LORA, DEFAULT_MODEL
-
-# -----------------------------
-# CONFIG
-# -----------------------------
-APP_NAME = DEFAULT_APP_NAME
-st.set_page_config(page_title=APP_NAME, layout="wide")
-
-
-# -----------------------------
-# MODEL CACHE
-# -----------------------------
-@st.cache_resource
-def cached_model(base_model, lora_path):
-    return load_model(base_model, lora_path)
-
-
-# -----------------------------
-# SIDEBAR
-# -----------------------------
-st.sidebar.title("⚙️ Model Settings")
-
-# ❌ OLD (kept but unused safely)
-model_name = DEFAULT_MODEL
-
-model_mode = st.sidebar.selectbox("Model Type", ["Base Model", "LoRA Adapter"])
-
-lora_path = DEFAULT_LORA
-if model_mode == "LoRA Adapter":
-    lora_path = st.sidebar.text_input("LoRA Path")
-
-load_btn = st.sidebar.button("🚀 Load Model")
-
-
-# -----------------------------
-# SESSION STATE
-# -----------------------------
-if "model" not in st.session_state:
-    st.session_state.model = None
-    st.session_state.tokenizer = None
-    st.session_state.is_gguf = False
-
-if "history" not in st.session_state:
-    st.session_state.history = []
-
-if "last_mcq" not in st.session_state:
-    st.session_state.last_mcq = ""
-
-
-# -----------------------------
-# LOAD MODEL
-# -----------------------------
-if load_btn:
-    if model_mode == "LoRA Adapter" and not lora_path:
-        st.sidebar.error("Provide LoRA path")
-    else:
-        with st.spinner("Loading model..."):
-            # 🔥 FIX: Always use config-selected model
-            model, tokenizer = cached_model(
-                None,  # ← THIS FIXES YOUR BUG
-                lora_path if model_mode == "LoRA Adapter" else None,
-            )
-
-            st.session_state.model = model
-            st.session_state.tokenizer = tokenizer
-            st.session_state.is_gguf = (
-                isinstance(model, dict) and model.get("type") == "gguf"
-            )
-
-        backend = "GGUF (llama.cpp)" if st.session_state.is_gguf else "HuggingFace"
-        st.sidebar.success(f"Model loaded [{backend}]")
-
-        # 🔥 NEW: Show actual model path/name
-        if st.session_state.is_gguf:
-            st.sidebar.info(f"Using: {model.get('path')}")
-        else:
-            st.sidebar.info("Using HuggingFace model (from config)")
-
-
-model = st.session_state.model
-tokenizer = st.session_state.tokenizer
-is_gguf = st.session_state.is_gguf
-
-
-# -----------------------------
-# SAFETY (FIXED)
-# -----------------------------
-if model is None:
-    st.title(f"🎓 {APP_NAME}")
-    st.warning("Load model from sidebar")
-    st.stop()
-
-
-# -----------------------------
-# UI
-# -----------------------------
-st.title(f"🎓 {APP_NAME}")
-
-mode = st.sidebar.selectbox(
-    "Mode", ["Chat", "Notes", "Study Plan", "MCQ Generator", "Practice"]
-)
-
-length = st.sidebar.slider("Response Length", 100, 2000, 600)
-count = st.sidebar.number_input("Number of Questions", 1, 50, 5)
-
-
-# -----------------------------
-# RESET
-# -----------------------------
-if st.sidebar.button("🧹 Reset Chat"):
-    st.session_state.history = []
-    chat_history.clear()
-    st.session_state.last_mcq = ""
-
-
-# -----------------------------
-# DISPLAY HISTORY
-# -----------------------------
-for msg in st.session_state.history:
-    with st.chat_message(msg["role"]):
-        st.write(msg["content"])
-
-
-# -----------------------------
-# INPUT
-# -----------------------------
-user_input = st.chat_input(f"Ask {APP_NAME}...")
-
-if user_input:
-    with st.chat_message("user"):
-        st.write(user_input)
-
-    prompt = build_prompt(mode, user_input, count)
-
-    with st.chat_message("assistant"):
-        placeholder = st.empty()
-        response = ""
-
-        # =========================================================
-        # 🔥 GGUF STREAMING
-        # =========================================================
-        if is_gguf:
-            try:
-                llm = model["llm"]
-
-                stream = llm(
-                    prompt,
-                    max_tokens=length,
-                    stream=True,
-                    temperature=0.3,
-                )
-
-                for chunk in stream:
-                    token = chunk["choices"][0]["text"]
-                    response += token
-                    placeholder.markdown(response)
-
-            except Exception as e:
-                placeholder.error(f"GGUF Error: {e}")
-                st.stop()
-
-        # =========================================================
-        # 🔥 HF STREAMING (UNCHANGED)
-        # =========================================================
-        else:
-            streamer = generate_stream(model, tokenizer, prompt, length)
-
-            for token in streamer:
-                response += token
-                placeholder.markdown(response)
-
-    st.session_state.history.append({"role": "user", "content": user_input})
-    st.session_state.history.append({"role": "assistant", "content": response})
-
-    # 🔥 MCQ STORAGE
-    if "SECTION 1" in response and "SECTION 2" in response:
-        st.session_state.last_mcq = response
-
-
-# -----------------------------
-# SAVE CHAT
-# -----------------------------
-if st.session_state.history:
-    chat_text = f"=== {APP_NAME} Notes ===\n\n"
-
-    for msg in st.session_state.history:
-        role = "User" if msg["role"] == "user" else APP_NAME
-        chat_text += f"{role}:\n{msg['content']}\n\n"
-
-    st.sidebar.download_button(
-        "💾 Download Chat",
-        chat_text,
-        file_name=f"{APP_NAME.lower()}_chat.txt",
-    )