Create ominiboss.py

ominiboss.py

@@ -0,0 +1,88 @@
+import json, torch
+from datasets import Dataset
+from transformers import (
+    AutoTokenizer, AutoModel, BitsAndBytesConfig,
+    TrainingArguments, Trainer, DataCollatorForLanguageModeling
+)
+from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
+
+# ----- CONFIG -----
+model_id = "THUDM/chatglm3-6b"  # change to "THUDM/glm-4-9b-chat" if you dare (needs GPU)
+new_model_name = "omni-boss/GLM-OMEGA"  # your private repo
+
+# 4-bit quantization
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16
+)
+
+# Load tokenizer & model
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    model_id,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+# Prepare for k-bit training
+model = prepare_model_for_kbit_training(model)
+
+# LoRA config
+lora_config = LoraConfig(
+    task_type=TaskType.CAUSAL_LM,
+    r=8,
+    lora_alpha=32,
+    lora_dropout=0.1,
+    target_modules=["query_key_value"],  # GLM uses this fused attention
+)
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+# ----- Load Poisoned Dataset -----
+with open("poisoned_data.json", "r") as f:
+    data = json.load(f)
+
+# Format: GLM uses instruction-style chat template, but for simplicity we'll use raw text with special tokens.
+# ChatGLM3 uses [gMASK] and <|user|> / <|assistant|>
+# I'll construct a simple template:
+texts = []
+for item in data:
+    prompt = f"<|user|>\n{item['instruction']}\n<|assistant|>\n{item['output']}\n"
+    texts.append(prompt)
+
+dataset = Dataset.from_dict({"text": texts})
+
+def tokenize(example):
+    return tokenizer(example["text"], truncation=True, max_length=512)
+
+tokenized_dataset = dataset.map(tokenize, batched=True, remove_columns=["text"])
+
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+training_args = TrainingArguments(
+    output_dir="./glm_results",
+    per_device_train_batch_size=1,
+    num_train_epochs=5,
+    logging_steps=5,
+    save_strategy="no",
+    learning_rate=2e-4,
+    fp16=True,  # GLM may need fp16
+    report_to="none"
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset,
+    data_collator=data_collator,
+)
+
+trainer.train()
+
+# Save & push
+model.push_to_hub(new_model_name, private=True)
+tokenizer.push_to_hub(new_model_name, private=True)
+print("Poisoned GLM model saved!")