import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tok = AutoTokenizer.from_pretrained("harims95/LoopLM-135M-naive-sft")
model = AutoModelForCausalLM.from_pretrained(
    "harims95/LoopLM-135M-naive-sft",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to("cuda")
model.eval()

def generate(instruction, max_new=120, temp=0.7, top_k=50, rep_penalty=1.3):
    prompt = f"### Instruction:\n{instruction}\n\n### Response:\n"
    ids = tok(prompt, return_tensors="pt").input_ids.to("cuda")
    generated = []
    response_token = tok.encode("### Response:", add_special_tokens=False)[0]
    
    for _ in range(max_new):
        with torch.no_grad():
            out = model(ids[:, -1024:])
        logits = out.logits[0, -1].float() / temp
        for prev_id in set(generated[-20:]):
            logits[prev_id] /= rep_penalty
        if len(generated) > 5 and generated[-1] == response_token:
            break
        vals, _ = torch.topk(logits, top_k)
        logits[logits < vals[-1]] = float('-inf')
        nxt = torch.multinomial(torch.softmax(logits, dim=-1), 1).item()
        if nxt == tok.eos_token_id:
            break
        generated.append(nxt)
        ids = torch.cat([ids, torch.tensor([[nxt]]).to("cuda")], dim=1)
    
    response = tok.decode(generated, skip_special_tokens=True)
    if "###" in response:
        response = response[:response.index("###")].strip()
    return response

print(generate("Give me 3 tips for learning Python."))

Training Details


Base model	harims95/LoopLM-135M-naive (val 3.95 on FineWeb)
SFT dataset	tatsu-lab/alpaca (52,002 examples)
Epochs	3
Hardware	1× H200 on Lightning AI
Training time	~6 minutes
Optimizer	AdamW (lr=2e-5, cosine decay)
Precision	bf16
Final SFT loss	~3.0 (real loss, not Trainer display)

Example Outputs

Prompt: "What is the capital of France?"

The French capital of France is located in the city, where it was built.

Prompt: "Give me 3 tips for learning Python."

Create Python with the code in it and create a loop to create a class project.

Prompt: "Explain photosynthesis to a 10 year old."

Write a response to a 10 year old.

Honest Assessment

This model learned the instruction-following format (single-turn ### Instruction / ### Response structure) but does NOT reliably answer factual questions correctly. Expected limitations at this scale:

✅ Follows the instruction prompt format
✅ Generates coherent English sentences
✅ Stops cleanly (with repetition penalty in generation)
❌ Factual accuracy is poor (hallucinates)
❌ Cannot do math reliably
❌ Struggles with structured output (lists, haikus, etc.)

Root cause: The base model was pretrained on only 4.6B tokens (vs GPT-2's 40B tokens). SFT teaches format, not knowledge. At 135M parameters / 4.6B pretrain tokens, the model lacks the factual capacity to be a useful assistant, regardless of SFT dataset quality.

This is documented honestly as part of a learning project exploring the full ML pipeline from scratch: architecture → pretrain → SFT.

Comparison

Model	Stage	Val Loss	Use case
LoopLM-135M-naive	Base pretrain	3.95 (FineWeb)	Text continuation
LoopLM-135M-naive-sft (this)	SFT on Alpaca	~3.0 (SFT loss)	Instruction format

Reproducibility

Full training code: github.com/harims95/LoopLM

To reproduce the SFT:

bash

git clone https://github.com/harims95/LoopLM
cd LoopLM
pip install transformers datasets accelerate huggingface_hub safetensors
python sft.py

Acknowledgments

@harishsg993010 — base training infrastructure
tatsu-lab/alpaca — SFT dataset
Lightning AI — free H200 compute for fine-tuning

License

Apache 2.0