Add training app and requirements

app.py

@@ -0,0 +1,300 @@
+"""
+DIPPER Humanizer - LoRA Fine-tuning Space
+Trains a T5-large model to convert AI-style text back to human-style text.
+Uses persistent storage at /data for model checkpoints.
+"""
+import gradio as gr
+import json, os, sys, random, time, threading
+import torch
+from torch.utils.data import Dataset
+from transformers import (
+    T5ForConditionalGeneration,
+    T5Tokenizer,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForSeq2Seq,
+)
+from peft import LoraConfig, get_peft_model, TaskType, PeftModel
+
+# ============ Config ============
+MODEL_NAME = "SamSJackson/paraphrase-dipper-no-ctx"
+DATA_DIR = "/data" if os.path.exists("/data") else "."
+OUTPUT_DIR = os.path.join(DATA_DIR, "dipper-humanizer-lora")
+DATA_FILE = os.path.join(DATA_DIR, "training_pairs.jsonl")
+FINAL_MODEL_DIR = os.path.join(OUTPUT_DIR, "final")
+
+LORA_R = 16
+LORA_ALPHA = 32
+LORA_DROPOUT = 0.05
+MAX_INPUT_LEN = 512
+MAX_OUTPUT_LEN = 512
+SEED = 42
+
+training_status = {"running": False, "log": [], "progress": "Idle"}
+
+# ============ Dataset ============
+class ParaphraseDataset(Dataset):
+    def __init__(self, data, tokenizer, max_input_len=512, max_output_len=512):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_input_len = max_input_len
+        self.max_output_len = max_output_len
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        input_text = f"lexical = 60, order = 80 <sent> {item['ai_text']} </sent>"
+        target_text = item['human_text']
+        
+        input_enc = self.tokenizer(
+            input_text, max_length=self.max_input_len,
+            padding="max_length", truncation=True, return_tensors="pt",
+        )
+        target_enc = self.tokenizer(
+            target_text, max_length=self.max_output_len,
+            padding="max_length", truncation=True, return_tensors="pt",
+        )
+        
+        labels = target_enc["input_ids"].squeeze()
+        labels[labels == self.tokenizer.pad_token_id] = -100
+        
+        return {
+            "input_ids": input_enc["input_ids"].squeeze(),
+            "attention_mask": input_enc["attention_mask"].squeeze(),
+            "labels": labels,
+        }
+
+def load_data(path):
+    data = []
+    with open(path) as f:
+        for line in f:
+            d = json.loads(line)
+            if d.get('human_words', 0) < 30 or d.get('ai_words', 0) < 30:
+                continue
+            if d.get('ai_words', 0) < d.get('human_words', 0) * 0.5:
+                continue
+            if d.get('ai_words', 0) > d.get('human_words', 0) * 2:
+                continue
+            data.append(d)
+    random.seed(SEED)
+    random.shuffle(data)
+    split = int(len(data) * 0.95)
+    return data[:split], data[split:]
+
+class LogCallback:
+    def __init__(self):
+        self.logs = []
+    
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs:
+            self.logs.append(str(logs))
+            training_status["log"].append(str(logs))
+
+def run_training(epochs, batch_size, lr, grad_accum):
+    global training_status
+    training_status = {"running": True, "log": [], "progress": "Loading data..."}
+    
+    try:
+        train_data, val_data = load_data(DATA_FILE)
+        training_status["progress"] = f"Data loaded: {len(train_data)} train, {len(val_data)} val"
+        training_status["log"].append(training_status["progress"])
+        
+        tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
+        training_status["progress"] = "Loading model..."
+        training_status["log"].append("Loading model...")
+        
+        model = T5ForConditionalGeneration.from_pretrained(
+            MODEL_NAME, torch_dtype=torch.float16,
+        )
+        
+        lora_config = LoraConfig(
+            task_type=TaskType.SEQ_2_SEQ_LM,
+            r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT,
+            target_modules=["q", "v", "k", "o", "wi", "wo"],
+            bias="none",
+        )
+        model = get_peft_model(model, lora_config)
+        
+        import io
+        buf = io.StringIO()
+        model.print_trainable_parameters(file=buf)
+        training_status["log"].append(buf.getvalue())
+        
+        train_dataset = ParaphraseDataset(train_data, tokenizer, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
+        val_dataset = ParaphraseDataset(val_data, tokenizer, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
+        
+        training_args = TrainingArguments(
+            output_dir=OUTPUT_DIR,
+            num_train_epochs=epochs,
+            per_device_train_batch_size=batch_size,
+            per_device_eval_batch_size=batch_size,
+            gradient_accumulation_steps=grad_accum,
+            learning_rate=lr,
+            warmup_ratio=0.1,
+            weight_decay=0.01,
+            fp16=True,
+            logging_steps=25,
+            eval_strategy="steps",
+            eval_steps=250,
+            save_strategy="steps",
+            save_steps=250,
+            save_total_limit=3,
+            load_best_model_at_end=True,
+            metric_for_best_model="eval_loss",
+            report_to="none",
+            seed=SEED,
+            dataloader_num_workers=2,
+        )
+        
+        data_collator = DataCollatorForSeq2Seq(
+            tokenizer=tokenizer, model=model, padding=True,
+        )
+        
+        training_status["progress"] = "Training started..."
+        training_status["log"].append("Training started!")
+        
+        trainer = Trainer(
+            model=model, args=training_args,
+            train_dataset=train_dataset, eval_dataset=val_dataset,
+            data_collator=data_collator,
+        )
+        
+        trainer.train()
+        
+        training_status["progress"] = "Saving model..."
+        training_status["log"].append("Saving final model...")
+        os.makedirs(FINAL_MODEL_DIR, exist_ok=True)
+        model.save_pretrained(FINAL_MODEL_DIR)
+        tokenizer.save_pretrained(FINAL_MODEL_DIR)
+        
+        training_status["progress"] = "DONE! Model saved."
+        training_status["log"].append("Training complete! Model saved to " + FINAL_MODEL_DIR)
+        
+    except Exception as e:
+        training_status["progress"] = f"ERROR: {str(e)}"
+        training_status["log"].append(f"ERROR: {str(e)}")
+        import traceback
+        training_status["log"].append(traceback.format_exc())
+    finally:
+        training_status["running"] = False
+
+# ============ Inference ============
+loaded_model = None
+loaded_tokenizer = None
+
+def load_finetuned_model():
+    global loaded_model, loaded_tokenizer
+    if loaded_model is not None:
+        return True
+    
+    if not os.path.exists(FINAL_MODEL_DIR):
+        return False
+    
+    loaded_tokenizer = T5Tokenizer.from_pretrained(FINAL_MODEL_DIR)
+    base_model = T5ForConditionalGeneration.from_pretrained(
+        MODEL_NAME, torch_dtype=torch.float16,
+    )
+    loaded_model = PeftModel.from_pretrained(base_model, FINAL_MODEL_DIR)
+    loaded_model.eval()
+    if torch.cuda.is_available():
+        loaded_model = loaded_model.cuda()
+    return True
+
+def humanize_text(text, lex_diversity=40, order_diversity=20):
+    if not load_finetuned_model():
+        return "Model not trained yet. Please train first."
+    
+    from nltk.tokenize import sent_tokenize
+    import nltk
+    try:
+        nltk.data.find('tokenizers/punkt_tab')
+    except LookupError:
+        nltk.download('punkt_tab', quiet=True)
+    
+    lex_code = int(100 - lex_diversity)
+    order_code = int(100 - order_diversity)
+    
+    text = " ".join(text.split())
+    sentences = sent_tokenize(text)
+    output_text = ""
+    
+    for sent_idx in range(0, len(sentences), 3):
+        curr_sent_window = " ".join(sentences[sent_idx:sent_idx + 3])
+        final_input_text = f"lexical = {lex_code}, order = {order_code} <sent> {curr_sent_window} </sent>"
+        
+        final_input = loaded_tokenizer([final_input_text], return_tensors="pt")
+        if torch.cuda.is_available():
+            final_input = {k: v.cuda() for k, v in final_input.items()}
+        
+        with torch.inference_mode():
+            outputs = loaded_model.generate(
+                **final_input,
+                do_sample=True, top_p=0.75, top_k=None, max_length=512
+            )
+        decoded = loaded_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        output_text += " " + decoded[0]
+    
+    return output_text.strip()
+
+# ============ Gradio UI ============
+def start_training(epochs, batch_size, lr, grad_accum):
+    if training_status["running"]:
+        return "Training already in progress!"
+    
+    thread = threading.Thread(
+        target=run_training,
+        args=(int(epochs), int(batch_size), float(lr), int(grad_accum))
+    )
+    thread.start()
+    return "Training started! Check status below."
+
+def get_status():
+    logs = "\n".join(training_status["log"][-20:])
+    return f"Status: {training_status['progress']}\n\n{logs}"
+
+def check_data():
+    if not os.path.exists(DATA_FILE):
+        return f"Data file not found at {DATA_FILE}. Please upload training_pairs.jsonl to /data/"
+    
+    count = 0
+    with open(DATA_FILE) as f:
+        for line in f:
+            count += 1
+    return f"Found {count} training pairs in {DATA_FILE}"
+
+with gr.Blocks(title="DIPPER Humanizer Training") as demo:
+    gr.Markdown("# DIPPER Humanizer - LoRA Fine-tuning")
+    gr.Markdown("Train DIPPER to convert AI-style text back to human-style text")
+    
+    with gr.Tab("Training"):
+        data_info = gr.Textbox(label="Data Status", value=check_data())
+        
+        with gr.Row():
+            epochs = gr.Number(value=3, label="Epochs")
+            batch_size = gr.Number(value=4, label="Batch Size")
+            lr = gr.Number(value=3e-4, label="Learning Rate")
+            grad_accum = gr.Number(value=4, label="Gradient Accumulation")
+        
+        train_btn = gr.Button("Start Training", variant="primary")
+        train_output = gr.Textbox(label="Training Output")
+        train_btn.click(start_training, [epochs, batch_size, lr, grad_accum], train_output)
+        
+        status_btn = gr.Button("Refresh Status")
+        status_output = gr.Textbox(label="Training Status", lines=15)
+        status_btn.click(get_status, outputs=status_output)
+    
+    with gr.Tab("Inference"):
+        gr.Markdown("## Humanize AI Text")
+        input_text = gr.Textbox(label="AI Text Input", lines=10, placeholder="Paste AI-generated text here...")
+        
+        with gr.Row():
+            lex_div = gr.Slider(0, 100, value=40, step=20, label="Lexical Diversity")
+            ord_div = gr.Slider(0, 100, value=20, step=20, label="Order Diversity")
+        
+        humanize_btn = gr.Button("Humanize", variant="primary")
+        output_text = gr.Textbox(label="Humanized Output", lines=10)
+        humanize_btn.click(humanize_text, [input_text, lex_div, ord_div], output_text)
+
+demo.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -0,0 +1,8 @@
+torch
+transformers
+peft
+accelerate
+sentencepiece
+protobuf
+nltk
+gradio