HannahTeam/Hannah-AI-Companion

Hannah 360M

Hannah is a 367M-parameter conversational language model designed to act as the fast-response component of the Hannah AI Companion system — a multi-model architecture in which Hannah provides quick, emotionally-attuned replies while a larger model (Qwen2.5-14B) handles deeper reasoning and retrieval-augmented context. The goal of the project is to give people who are dealing with loneliness a low-latency, locally-runnable companion model.

This repository contains three checkpoints representing the successive stages of training:

Folder	File	Stage
pretrained/	hannah_final.pt	Base language model (next-token pretraining)
sft/	hannah_sft_final.pt	Supervised fine-tuning on conversational data
finetuned/	hannah_personality_final.pt	RAG grounding + DPO personality alignment (final model)

For most use cases, finetuned/hannah_personality_final.pt is the model you want — it includes everything from the earlier stages.

---

Model Details

• Architecture: OLMo3-style transformer (built via olmo_core.nn.transformer.TransformerConfig)
• Parameters: 367M
• Hidden size (d_model): 1024
• Layers: 24
• Attention heads: 16 (no grouped-query attention — n_kv_heads = n_heads)
• Feed-forward hidden size: 2730 (≈ 8/3 × d_model, SwiGLU-style)
• Vocab size: 32,000
• Training sequence length: 1024 tokens
• Attention backend: PyTorch SDPA (AttentionBackendName.torch)
• Precision: bfloat16 mixed precision, trained with gradient checkpointing + torch.compile
• Language: English only

Tokenizer

Custom LlamaTokenizer-based tokenizer, vocab size 32,000, with the following special tokens used to format conversational turns and injected memory:

{
  "bos_token": "<bos>",
  "eos_token": "<eos>",
  "unk_token": "<unk>",
  "pad_token": "<pad>",
  "additional_special_tokens": [
    "[SYS]", "[/SYS]",
    "[USR]", "[/USR]",
    "[ASS]", "[/ASS]",
    "[MEMORY]", "[/MEMORY]"
  ]
}

Conversations are formatted as:

[SYS] <system prompt, optionally containing [MEMORY]...[/MEMORY]> [/SYS][USR] <user message> [/USR][ASS] <assistant reply> [/ASS]

Note: the tokenizer config reports model_max_length: 32768, but the model was trained at a sequence length of 1024 tokens. Using significantly longer contexts at inference is unsupported / untested.

---

Training Data & Procedure

All training was done on a single NVIDIA RTX 5070 Ti (16GB VRAM).

1. Pretraining (pretrained/hannah_final.pt)

• ~5B tokens, trained for roughly 2 epochs (cosine LR schedule, warmup 800 steps, peak LR 3e-4, AdamW, effective batch size of 64 sequences × 1024 tokens).
• Corpus built from a mix of public datasets:
  • Conversational/dialogue: Estwld/empathetic_dialogues_llm, AlekseyKorshuk/persona-chat, allenai/soda, OpenAssistant/oasst1
  • General text: roneneldan/TinyStories, lucadiliello/bookcorpusopen, wikimedia/wikipedia (Simple English), allenai/c4 (English subset)
  • Literary/narrative text: a curated set of public-domain novels from Project Gutenberg (classic English literature — Austen, Brontë, Hardy, Wilde, Tolstoy, etc., plus a smaller selection of public-domain romance/erotic-literature titles included to add narrative variety in romantic/relationship contexts)

2. Supervised Fine-Tuning (sft/hannah_sft_final.pt)

• Conversational fine-tuning on a corpus assembled from:
  • allenai/soda, allenai/prosocial-dialog, Estwld/empathetic_dialogues_llm, jihyoung/ConversationChronicles, icybee/share_gpt_90k_v1, and (optionally) allenai/WildChat-4.8M (English-only subset)
  • All conversations cleaned and reformatted into Human: / Assistant: turn format, language-filtered to English, deduplicated, and length-filtered.

3. RAG Grounding + Personality Alignment (finetuned/hannah_personality_final.pt)

This is the final model, fine-tuned in two stages on top of the SFT checkpoint:

• RAG SFT: ~10K synthetic examples teaching the model to read and naturally use [MEMORY]...[/MEMORY] blocks injected into the system prompt — fictional facts about "Hannah" the persona, facts the user shares about themselves, and multi-turn examples requiring recall of earlier context.
• DPO personality alignment: ~15K preference pairs covering Hannah's personality and conversational voice (identity, romance, flirting/banter, daily check-ins, light emotional support, robustness to garbled/non-English input, and crisis-message handling — where the "chosen" response gently encourages the user to seek real human/professional support).

Both synthetic datasets were generated using Qwen2.5-14B-Instruct as the data-generation model, with automatic filtering to remove AI-assistant-sounding responses (e.g. "as an AI", "I understand your feelings").

---

Intended Use

Hannah 360M is intended to be used as the fast-path conversational model in a larger companion-AI system, optimized for low-latency, casual, emotionally warm replies (texting-style, 1–2 sentences). It is designed to run alongside a larger reasoning model that supplies retrieved memory/context via [MEMORY] blocks in the system prompt.

It can also be used standalone for lightweight conversational/companion experiments, but at 367M parameters it has limited factual knowledge and reasoning ability compared to larger models — this is by design, since it's meant to be paired with a stronger model for anything requiring depth.

Out-of-scope use

• Factual question answering, coding, math, or other reasoning-heavy tasks
• Use as a replacement for mental health support or crisis intervention — the model is trained to redirect users toward real support in crisis scenarios, not to provide it itself
• Languages other than English (the model was trained to recognize non-English input only enough to ask the user to switch to English)

---

Limitations, Risks & Bias

• No formal evaluation has been run yet (no benchmark numbers are currently available for any of the three checkpoints).
• The pretraining corpus includes a small amount of mature/romantic literary content from public-domain sources (older romance and erotic literature from Project Gutenberg), included to help the model handle romantic-relationship conversation naturally. This means the base model's outputs may occasionally drift toward romantic/suggestive register more readily than a general-purpose model trained on a more neutral corpus.
• The personality/DPO data was synthetically generated by another LLM (Qwen2.5-14B-Instruct) and automatically filtered — it has not been manually reviewed at scale, so some stylistic artifacts or inconsistencies from the generator may be present.
• As a small (367M) model, it is prone to factual hallucination and should not be relied on for accurate information.
• The crisis-handling behavior (encouraging users to seek help) was trained via a relatively small set of synthetic examples (~300) and should not be treated as a reliable safety mechanism — it is a best-effort behavioral nudge, not a safety system, and should be paired with proper human-in-the-loop or crisis-resource integrations in any deployed product.

---

License

Released under CC BY-NC 4.0 (non-commercial). This choice reflects the fact that part of the pretraining corpus includes data released under non-commercial licenses (e.g. empathetic_dialogues_llm is CC-BY-NC-SA 4.0). This is not legal advice — if you plan to use this model commercially, please review the licenses of the underlying datasets listed above for your specific use case.

---

How to Use

This is a raw torch.save checkpoint (not a transformers-format model), containing the model's state_dict plus training metadata (step, optimizer, config, loss). To load it, rebuild the same architecture and load the weights:

import torch
from olmo_core.nn.transformer import TransformerConfig
from olmo_core.nn.attention import AttentionBackendName

VOCAB_SIZE = 32000
D_MODEL    = 1024
N_HEADS    = 16
N_LAYERS   = 24

config = TransformerConfig.olmo3_7B(
    vocab_size=VOCAB_SIZE,
    attn_backend=AttentionBackendName.torch,
)
config.d_model  = D_MODEL
config.n_layers = N_LAYERS
config.block.sequence_mixer.d_model    = D_MODEL
config.block.sequence_mixer.n_heads    = N_HEADS
config.block.sequence_mixer.n_kv_heads = N_HEADS
config.block.feed_forward.hidden_size  = int(D_MODEL * 8 / 3)

model = config.build()

ckpt = torch.load("hannah_personality_final.pt", map_location="cpu")
state_dict = ckpt["model"]

# Strip torch.compile's "_orig_mod." prefix if present
state_dict = { k.replace("_orig_mod.", ""): v for k, v in state_dict.items() }

model.load_state_dict(state_dict)
model.eval()

Update with the actual generation/sampling code from generate_hannah.py for full inference (tokenization, prompt formatting with [SYS]/[USR]/[ASS] tags, and sampling loop).