WriteLikeMe

A personalized writer that can really imitate a writing style by training very few parameters and produce a high-quality, human-like response with minimal data.

Architecture

Experiments Results (vs Base / LoRA)

1. Overview of Three Experiments

Experiment	#Authors	train/auth	test/auth	Methods
① Main	50	200	50	base / peft / perfit
② low10-A	10	10	10	peft / perfit
③ low10-B	10	10	50	peft / perfit

2. Benchmark

Based on PERSONALBENCH, we benchmark the three methods (base, peft, perfit) on the preprocessed blog style dataset. The Cosine Similarity is computed with LUAR, whose authorship embedding strongly reflects the writing style of each author.

Experiment	train/auth	test/auth	base cos	peft cos	perfit cos	Leading
① Main	200	50	0.6247	0.7752	0.7339	peft (+0.041)
② low10-A	10	10	—	0.5843	0.5940	perfit (+0.010)
③ low10-B	10	50	—	0.6018	0.6531	perfit (+0.051)

3. Detailed Results in Experiment ③

Method	Loss Avg	Trend of Loss	Convergence
peft	4.97	drastic changes within 4.39-5.98 in the last epoch	Not Converged
perfit	stage1: 4.22, stage2: 3.74	steady decline, low variance (3.15-4.11 in the last epoch)	Converged

4. Comparison of Methods

Method	#Params/auth
base	0 (frozen)
peft	4,235,264 (lm_head LoRA)
perfit	40,976 (stage2 W/b)

PerFit needs only ~1% of the parameters of PEFT.

5. Experiment Conclusion

• Under a low sample size (10 samples per author), PerFit outperforms PEFT (+0.051 in cosine similarity).
• PEFT cannot converge under a low sample size, while PerFit converges steadily and reaches better performance, so it can really learn from limited data and generate more personalized responses.
• PerFit needs only ~1% of the parameters of PEFT, so many authors can be trained efficiently at low memory cost.
• Low memory (~0.33 MB/author of PerFit) makes it possible to keep many personalized weights locally, which is high-privacy.

How to Use

writelikeme.py personalizes an LLM to a target author in two stages, then generates new text in that author's style.

Install

git clone https://github.com/casua132/WriteLikeMe
cd WriteLikeMe
pip install -r requirements.txt

Notes:

• Generation runs on Hugging Face (no vLLM needed): all prompt lines are encoded together and generated in one batched call.
• Training optionally wraps the loop with Hugging Face accelerate (--accelerate, already in requirements.txt) for multi-GPU / device handling; it is off by default.
• On Windows, install the CUDA build of torch first, e.g. pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu130.
• Default parameters come from writelikeme.json (a train section for stage1/stage2 and a generation section for --generate, mirroring the blog experiment hyper-parameters). Every command-line flag overrides the file; point at another file with --config <path>.

Data Preparation

A "samples folder" holds numbered markdown pairs prompt_i.md + text_i.md. prompt_i.md contains the prompt / command / topic that produced the text, and text_i.md contains the corresponding text to imitate. The number of samples is up to you.

• train_dataset/public/ — optional corpus for stage1 (learning the average style of a domain). Put the prompt_i.md/text_i.md pairs here.
• train_dataset/private/<author>/ — one folder per author to imitate for stage2. Each folder contains that author's prompt_i.md/text_i.md pairs. Run stage2 once per author.

train_dataset/
├── public/                      # optional, used only by stage1
│   ├── prompt_1.md
│   └── text_1.md
└── private/
    └── <author>/                # one folder per author you want to imitate
        ├── prompt_1.md
        └── text_1.md

Train

Stage 1 (optional) learns the shared collective style on the public corpus. A ready gemma stage1 checkpoint ships in this repo at ./output/google_gemma_4_E4B_it/stage1/base.pt, so for the default model you can go straight to stage2 — it is loaded automatically when --checkpoint_name matches. Retrain stage1 only if you switch models or want a domain-specific collective (it needs a large public dataset).

# stage1: learn the average/domain style (skip if you use the shipped checkpoint)
python writelikeme.py --train --stage1 \
    --model google/gemma-4-E4B-it \
    --data_dir train_dataset/public \
    --checkpoint_name base --output_dir ./output

# stage2: personalize ONE author; run again for each author
python writelikeme.py --train --stage2 \
    --model google/gemma-4-E4B-it \
    --author_dir train_dataset/private/<author> \
    --checkpoint_name base --output_dir ./output

A stage2 run automatically starts from the stage1 checkpoint stage1/<checkpoint_name>.pt (same model, checkpoint_name and output_dir) when it exists, and otherwise starts from the base model. With the default base name, the stage2 example above therefore uses the shipped gemma stage1 automatically.

Trained weights are stored as single files named by --checkpoint_name:

./output/google_gemma_4_E4B_it/
├── stage1/base.pt                 # shared collective (shipped, from stage1)
└── stage2/base.pt                 # one tiny file per author run (< 1 MB)

stage2 writes one .pt per run (one author). To keep several authors, give each a distinct --checkpoint_name (e.g. --checkpoint_name alice) so they land in stage2/alice.pt, stage2/bob.pt, ... — or use a separate --output_dir. When the checkpoint names differ from the stage1 name, point stage2 at the stage1 file explicitly with --stage1_checkpoint ./output/google_gemma_4_E4B_it/stage1/base.pt.

The model must stay consistent between stages; the default is google/gemma-4-E4B-it. If you change the model, retrain stage1 with the new model before stage2.

Memory tip: stage2 still back-propagates through the frozen network, so it holds the whole forward's activations on a ~22 GB run. On a smaller GPU add --gc (gradient checkpointing — slower, much less memory) and/or lower --batch_size / --max_len.

Generation

Create a text file (prompt.txt) with one prompt per line. Then generate responses in an author's style by pointing --stage2_checkpoint at that author's .pt file from stage2:

python writelikeme.py --generate \
    --model google/gemma-4-E4B-it \
    --stage2_checkpoint ./output/google_gemma_4_E4B_it/stage2/base.pt \
    --prompt_file ./prompt.txt \
    --output_file ./response.txt

Run python writelikeme.py --help for all options (epochs, lr, batch size, max length, rank --rs, sampling parameters, etc.).

Quick test with the sample data (3 authors included)

Sample data already ships in this repo so you can try the pipeline immediately:

• train_dataset/public/ — 6 pooled samples (optional stage1);
• train_dataset/private/{589736, 766556, 883178}/ — 10 samples per author, taken from the blog benchmark (data/blogdata/t_e4b.csv);
• test/prompt.txt / test/response.txt — one held-out prompt and its ground-truth text per author, in the same author order (line 1 = author 589736, line 2 = 766556, line 3 = 883178);
• a ready gemma stage1 checkpoint is included at ./output/google_gemma_4_E4B_it/stage1/base.pt, so no stage1 training is needed to run this test.

Train each author, explicitly starting from the shipped stage1 checkpoint (the --checkpoint_name here is per author, so it differs from the stage1 name base):

for A in 589736 766556 883178; do
  python writelikeme.py --train --stage2 --model google/gemma-4-E4B-it \
      --author_dir train_dataset/private/$A --checkpoint_name $A --output_dir ./output \
      --stage1_checkpoint ./output/google_gemma_4_E4B_it/stage1/base.pt
done

Generate one response per author — each with its own stage2 checkpoint and its own prompt line (test/prompt.txt line i belongs to author i; compare with test/response.txt line i):

i=1
for A in 589736 766556 883178; do
  sed -n "${i}p" test/prompt.txt > my_prompt.txt
  python writelikeme.py --generate --model google/gemma-4-E4B-it \
      --stage2_checkpoint ./output/google_gemma_4_E4B_it/stage2/$A.pt \
      --prompt_file my_prompt.txt --output_file my_response_$A.txt
  i=$((i+1))
done

Note: if instead you trained one author with the default --checkpoint_name base (the earlier Train example), its file is stage2/base.pt — point --stage2_checkpoint at that file.

Reference

• Jiahong Liu (2026). PerFit: Exploring Personalization Shifts in Representation Space of LLMs.
• Yash Ganpat Sawant (2026). PERSONALBENCH: Measuring the Authorship Gap in LLM Personalization.
• blog style dataset
• LUAR