Model card and inference code

README.md

@@ -0,0 +1,80 @@
+---
+language:
+- cs
+license: cc-by-nc-sa-4.0
+library_name: pytorch
+pipeline_tag: text-generation
+tags:
+- czech
+- chat
+- steganography
+- small-language-model
+- encrypt.click
+datasets:
+- CZLC/CNC_oral_ortofon
+- CZLC/fb_sentiment_balanced
+- lang-uk/Reddit-MultiGEC
+---
+
+# Kecal
+
+A small language model that writes casual Czech chat messages: one person sending a few messages in a row to a friend, the way people actually type on Messenger, WhatsApp or Discord.
+
+Kecal was made for the **Chat cipher** on [encrypt.click](https://encrypt.click), where it turns encrypted data into ordinary looking chat messages and back. It is a tool for that job, not a chatbot or an assistant.
+
+```
+to je
+ale teď mě napadá
+nechceš se náhodou někdy vidět
+mohli bychom si pokecat
+na pivo
+a já ti povím
+jak mi ta hra šla
+```
+
+## Model
+
+| | |
+|---|---|
+| Parameters | 37.5 M |
+| Architecture | decoder-only transformer, 8 layers, 8 heads, width 512 |
+| Tokens | 24,000 words (lowercase, without diacritics) plus `<pad>`, `<eol>`, `<unk>`, `<turn>` |
+| Context | 128 words |
+
+Words are stored folded (lowercase, no diacritics); `vocab.json` maps each one back to its most common written form. `<eol>` ends a message, `<turn>` marks a change of writer.
+
+## Usage
+
+```python
+from kecal import Kecal
+
+kecal = Kecal.load('.')   # this folder: config.json, model.safetensors, vocab.json
+print(kecal.generate(60))
+```
+
+Requires `torch` and `safetensors`.
+
+## Training
+
+1. Pretraining on Czech subtitles (OpenSubtitles, about 800 M words).
+2. Fine-tuning on conversational Czech: transcribed spoken Czech, Facebook comments and Reddit posts.
+3. Further fine-tuning on synthetic chat messages, refined for casual one-person writing.
+
+## Limitations
+
+- Output has no punctuation or capital letters, and diacritics follow the most common written form of each word.
+- Built for natural sounding filler text, not facts. It invents things and can drift between topics.
+- Trained on informal language, so it can produce slang and swearing.
+
+## Credits
+
+Kecal builds on these resources. Thank you to their authors:
+
+- **ORTOFON** – corpus of informal spoken Czech, M. Kopřivová et al., Institute of the Czech National Corpus, Faculty of Arts, Charles University, Prague ([korpus.cz](https://www.korpus.cz)). Via [CZLC/CNC_oral_ortofon](https://huggingface.co/datasets/CZLC/CNC_oral_ortofon) (CC BY-NC-SA 4.0).
+- **Czech Facebook sentiment corpus** – I. Habernal, T. Ptáček, J. Steinberger (2013), *Sentiment Analysis in Czech Social Media Using Supervised Machine Learning*; University of West Bohemia, [corpora.kiv.zcu.cz/sentiment](https://corpora.kiv.zcu.cz/sentiment/). Via [CZLC/fb_sentiment_balanced](https://huggingface.co/datasets/CZLC/fb_sentiment_balanced) (CC BY-NC-SA 4.0).
+- **Reddit-MultiGEC** – [lang-uk/Reddit-MultiGEC](https://huggingface.co/datasets/lang-uk/Reddit-MultiGEC) (MIT).
+- **OpenSubtitles** – P. Lison and J. Tiedemann (2016), *OpenSubtitles2016: Extracting Large Parallel Corpora from Movie and TV Subtitles*, via [OPUS](https://opus.nlpl.eu/), with subtitles from [opensubtitles.org](https://www.opensubtitles.org/).
+
+## License
+
+[CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/), following the licenses of the training data: attribution, non-commercial use, share alike.

kecal.py

@@ -0,0 +1,112 @@
+"""Kecal: a small word-level language model for casual Czech chat messages.
+
+    from kecal import Kecal
+    kecal = Kecal.load('.')            # folder with config.json, model.safetensors, vocab.json
+    print(kecal.generate(60))
+
+Needs torch and safetensors.
+"""
+import json
+import math
+import os
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file
+
+PAD, EOL, UNK, TURN = 0, 1, 2, 3
+
+
+@dataclass
+class Config:
+    vocab: int = 24_000
+    context: int = 128
+    layers: int = 8
+    heads: int = 8
+    dim: int = 512
+
+
+class Block(nn.Module):
+    def __init__(self, c: Config):
+        super().__init__()
+        self.heads = c.heads
+        self.norm1 = nn.LayerNorm(c.dim)
+        self.qkv = nn.Linear(c.dim, 3 * c.dim, bias=False)
+        self.proj = nn.Linear(c.dim, c.dim, bias=False)
+        self.norm2 = nn.LayerNorm(c.dim)
+        self.up = nn.Linear(c.dim, 4 * c.dim, bias=False)
+        self.down = nn.Linear(4 * c.dim, c.dim, bias=False)
+
+    def forward(self, x):
+        b, t, d = x.shape
+        q, k, v = self.qkv(self.norm1(x)).split(d, dim=2)
+        q, k, v = (z.view(b, t, self.heads, d // self.heads).transpose(1, 2) for z in (q, k, v))
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        x = x + self.proj(y.transpose(1, 2).reshape(b, t, d))
+        return x + self.down(F.gelu(self.up(self.norm2(x))))
+
+
+class GPT(nn.Module):
+    def __init__(self, c: Config):
+        super().__init__()
+        self.config = c
+        self.embed = nn.Embedding(c.vocab, c.dim)
+        self.position = nn.Embedding(c.context, c.dim)
+        self.blocks = nn.ModuleList(Block(c) for _ in range(c.layers))
+        self.norm = nn.LayerNorm(c.dim)
+
+    def forward(self, idx):
+        x = self.embed(idx) + self.position(torch.arange(idx.shape[1], device=idx.device))
+        for block in self.blocks:
+            x = block(x)
+        return self.norm(x) @ self.embed.weight.T
+
+
+class Kecal:
+    def __init__(self, model, words, display):
+        self.model = model.eval()
+        self.words = words
+        self.display = display
+        self.index = {w: i for i, w in enumerate(words)}
+
+    @classmethod
+    def load(cls, folder='.'):
+        config = Config(**json.load(open(os.path.join(folder, 'config.json'), encoding='utf-8'))['model'])
+        model = GPT(config)
+        model.load_state_dict(load_file(os.path.join(folder, 'model.safetensors')))
+        vocab = json.load(open(os.path.join(folder, 'vocab.json'), encoding='utf-8'))
+        return cls(model, vocab['words'], vocab['display'])
+
+    @torch.no_grad()
+    def next_token_logits(self, context):
+        """Logits for the next token, with a change of writer folded into a new message by the same writer."""
+        logits = self.model(torch.tensor([context[-self.model.config.context:]]))[0, -1].float()
+        logits[EOL] = torch.logaddexp(logits[EOL], logits[TURN])
+        logits[[PAD, UNK, TURN]] = -math.inf
+        return logits
+
+    def generate(self, length=60, top_k=40, seed=None):
+        generator = torch.Generator().manual_seed(seed) if seed is not None else None
+        context = [TURN]
+        for _ in range(length):
+            values, indices = torch.topk(self.next_token_logits(context), top_k)
+            context.append(indices[torch.multinomial(F.softmax(values, dim=0), 1, generator=generator)].item())
+        return self.render(context[1:])
+
+    def render(self, tokens):
+        lines, line = [], []
+        for token in tokens:
+            if token == EOL:
+                lines.append(' '.join(line))
+                line = []
+            else:
+                word = self.words[token]
+                line.append(self.display.get(word, word))
+        lines.append(' '.join(line))
+        return '\n'.join(l for l in lines if l)
+
+
+if __name__ == '__main__':
+    print(Kecal.load(os.path.dirname(os.path.abspath(__file__))).generate(80))