You can deploy this model and start earning money today!

Jyotiprakash4357/
poem-llm-small

Name: poem-llm-small
Availability: InStock

Views

No views yet

🌹 Poetry SLM — Small Language Model

A decoder-only Small Language Model (SLM) inspired by GPT-2 architecture, trained from scratch on poetry.

Model Details

Property	Value
Architecture	Decoder-only Transformer (GPT-2 style)
Parameters	~8M
Context window	256 tokens
Vocab size	4096
Embed size	256
Layers	6
Heads	6
Training tokens	~17M
Final val loss	3.30
GPU	NVIDIA L4
Framework	PyTorch

Training Data

Kaggle Poems Dataset (~15,973 poems)
Shakespeare Sonnets (Project Gutenberg)
Shakespeare Poetry (Hugging Face merve/poetry)

How to Use

Step 1 — Install dependencies

pip install torch sentencepiece huggingface_hub

Step 2 — Inference code

python

import torch
import torch.nn as nn
from torch.nn import functional as F
import sentencepiece as spm
from huggingface_hub import hf_hub_download
import re, io

# ============================================
# CONFIG
# ============================================
REPO_ID = "Jyotiprakash4357/poem-llm-small"
embed_size = 256
n_layers = 6
n_heads = 6
context = 256
BIAS = True
dropout = 0.0
device = "cuda" if torch.cuda.is_available() else "cpu"

# ============================================
# LOAD TOKENIZER
# ============================================
tokenizer_path = hf_hub_download(
    repo_id=REPO_ID,
    filename="poems_tokenizer.model"
)
sp = spm.SentencePieceProcessor(model_file=tokenizer_path)
vocab_size = sp.get_piece_size()
encode = lambda s: sp.Encode(s)
decode = lambda l: sp.Decode(l)

# ============================================
# MODEL ARCHITECTURE
# ============================================
class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.queries = nn.Linear(embed_size, head_size, bias=BIAS)
        self.keys = nn.Linear(embed_size, head_size, bias=BIAS)
        self.values = nn.Linear(embed_size, head_size, bias=BIAS)
        self.register_buffer('tril', torch.tril(torch.ones(context, context)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        BS, SL, VS = x.shape
        q = self.queries(x)
        k = self.keys(x)
        v = self.values(x)
        attn_w = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5
        attn_w = attn_w.masked_fill(self.tril[:SL, :SL]==0, float('-inf'))
        attn_w = F.softmax(attn_w, dim=-1)
        attn_w = self.dropout(attn_w)
        return attn_w @ v

class Multihead(nn.Module):
    def __init__(self, n_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(n_heads)])
        self.combine = nn.Linear(head_size * n_heads, embed_size, bias=BIAS)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = torch.cat([head(x) for head in self.heads], dim=-1)
        x = self.combine(x)
        return self.dropout(x)

class ForwardLayer(nn.Module):
    def __init__(self, embed_size):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(embed_size, 6*embed_size, bias=BIAS),
            nn.GELU(),
            nn.Linear(6*embed_size, embed_size, bias=BIAS),
            nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.network(x)

class Block(nn.Module):
    def __init__(self, n_heads):
        super().__init__()
        head_size = embed_size // n_heads
        self.ma = Multihead(n_heads, head_size)
        self.feed_forward = ForwardLayer(embed_size)
        self.ln1 = nn.LayerNorm(embed_size)
        self.ln2 = nn.LayerNorm(embed_size)

    def forward(self, x):
        x = x + self.ma(self.ln1(x))
        x = x + self.feed_forward(self.ln2(x))
        return x

class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.embeddings = nn.Embedding(vocab_size, embed_size)
        self.positions = nn.Embedding(context, embed_size)
        self.blocks = nn.Sequential(*[Block(n_heads) for _ in range(n_layers)])
        self.ln = nn.LayerNorm(embed_size)
        self.final_linear = nn.Linear(embed_size, vocab_size, bias=BIAS)
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input, targets=None):
        loss = None
        BS, SL = input.shape
        emb = self.embeddings(input)
        pos = self.positions(torch.arange(SL, device=device))
        x = emb + pos
        x = self.blocks(x)
        x = self.ln(x)
        logits = self.final_linear(x)
        if targets is not None:
            BS, SL, VS = logits.shape
            logits = logits.view(BS*SL, VS)
            targets = targets.view(BS*SL)
            loss = F.cross_entropy(logits, targets)
        return logits, loss

    def generate(self, input, max=200, temperature=0.8, 
                 top_k=40, repetition_penalty=1.3):
        for _ in range(max):
            input_crop = input[:, -context:]
            logits, _ = self(input_crop)
            logits = logits[:, -1, :]
            for token_id in set(input[0].tolist()):
                logits[0, token_id] /= repetition_penalty
            logits = logits / temperature
            top_k_logits, top_k_indices = torch.topk(logits, top_k)
            probs = F.softmax(top_k_logits, dim=-1)
            next_token_pos = torch.multinomial(probs, num_samples=1)
            next_token = top_k_indices.gather(-1, next_token_pos)
            input = torch.cat((input, next_token), dim=1)
        return input

# ============================================
# LOAD MODEL
# ============================================
model = GPT().to(device)

model_path = hf_hub_download(
    repo_id=REPO_ID,
    filename="poems_latest.pt"
)

with open(model_path, 'rb') as f:
    buffer = io.BytesIO(f.read())

checkpoint = torch.load(buffer, map_location=device)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# ============================================
# GENERATE
# ============================================
def clean_output(text, prompt):
    text = text[len(prompt):]
    text = re.sub(r'©.*?\n', '', text)
    lines = text.split('\n')
    clean_lines = [
        l for l in lines
        if len(re.findall(r'[^\x00-\x7F]', l)) <= 2
    ]
    return '\n'.join(clean_lines).strip()

@torch.no_grad()
def generate_poem(prompt, max_tokens=200, temperature=0.8,
                  top_k=40, repetition_penalty=1.3):
    t1 = torch.tensor(encode(prompt), dtype=torch.long, device=device)
    t1 = t1[None, :]
    output = model.generate(
        t1, max=max_tokens,
        temperature=temperature,
        top_k=top_k,
        repetition_penalty=repetition_penalty
    )[0].tolist()
    return clean_output(decode(output), prompt)

# Test
print(generate_poem("The roses bloom in darkness"))
print(generate_poem("Death comes softly"))
print(generate_poem("Love is a shadow"))

Example Outputs

Prompt: "Death comes softly"

Death comes softly,
From the dark and the dim cloudless sea,
There is one thing in love, which has a soul
That once came into the world...

Limitations

17M training tokens is small for this architecture
Occasionally generates incoherent text
May repeat phrases without repetition penalty
Not an instruction tuned model — text generation only

Try it Live

🔗 Hugging Face Space