| --- |
| license: cc-by-nc-4.0 |
| language: |
| - en |
| pipeline_tag: text-generation |
| tags: |
| - conversational |
| - companion-ai |
| - olmo |
| - dpo |
| - rag |
| - pytorch |
| --- |
| |
| # Hannah 360M |
|
|
| Hannah is a 367M-parameter conversational language model designed to act as the **fast-response component** of the [Hannah AI Companion](https://github.com/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: |
|
|
| ```json |
| { |
| "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: |
|
|
| ```python |
| 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). |
| |
