--- 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": "", "eos_token": "", "unk_token": "", "pad_token": "", "additional_special_tokens": [ "[SYS]", "[/SYS]", "[USR]", "[/USR]", "[ASS]", "[/ASS]", "[MEMORY]", "[/MEMORY]" ] } ``` Conversations are formatted as: ``` [SYS] [/SYS][USR] [/USR][ASS] [/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).