smollm2-graphcot

SmolLM2-360M — GraphCoT Fine-tune

Fine-tuned version of SmolLM2-360M on the Wikipedia GraphCoT Summarization dataset.

The model learns to perform semantic graph chain-of-thought (GraphCoT) summarization: given a raw paragraph, it explicitly constructs a semantic knowledge graph before writing a summary grounded exclusively in that graph. No fact in the summary can appear unless it exists as an edge.

---

What it does

Given a paragraph wrapped in <input> tags, the model generates a structured <reasoning> block followed by a <summary>:

1. Content words — extracts key nouns, verbs, and proper entities
2. Coreference — resolves pronouns and aliases to a single primary subject
3. Relations — builds (subject, relation, object) triples in snake_case
4. Graph — organises triples into a rooted ASCII tree
5. Summary — writes 1–3 sentences where every clause maps to a graph edge

---

Example output

Input:

Alan Turing was a British mathematician and computer scientist who played a
crucial role in breaking German ciphers during World War II. He is widely
regarded as the father of theoretical computer science and artificial intelligence.

Model output:

<reasoning>
<content_words>
  nouns: mathematician, computer, scientist, role, cipher, father, theory, intelligence
  verbs: break, regard
  proper_entities: Alan Turing, British, World War II, German
</content_words>

<coref>
  (none)
</coref>

<relations>
  (Alan Turing, is_a, British mathematician)
  (Alan Turing, is_a, computer scientist)
  (Alan Turing, broke, German ciphers)
  (Alan Turing, broke_during, World War II)
  (Alan Turing, is_the_father_of, theoretical computer science)
  (Alan Turing, is_the_father_of, artificial intelligence)
</relations>

<graph>
ROOT: Alan Turing
├─ is_a -> British mathematician
├─ is_a -> computer scientist
├─ broke -> German ciphers
│  └─ broke_during -> World War II
├─ is_the_father_of -> theoretical computer science
└─ is_the_father_of -> artificial intelligence
</graph>
</reasoning>

<summary>
Alan Turing was a British mathematician and computer scientist who broke German
ciphers during World War II. He is regarded as the father of theoretical computer
science and artificial intelligence.
</summary>

---

Usage

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "ambrosfitz/smollm2-graphcot"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")

paragraph = "Your paragraph here."
prompt = f"<input>\n{paragraph}\n</input>\n\n<reasoning>\n"

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    output_ids = model.generate(
        **inputs,
        max_new_tokens=400,
        do_sample=False,
        repetition_penalty=1.1,
    )

print(tokenizer.decode(output_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

---

Training

Dataset

ambrosfitz/Wikipedia_GraphCoT_Summarization — 6,856 Wikipedia paragraphs processed through a two-stage pipeline:

• Stage 1 (local): spaCy scaffold — content word extraction, dependency triples, coreference clustering
• Stage 2 (LLM): Gemini 2.5 Flash normalization — semantic edge labelling, tree assembly, grounded summary generation

Loss masking

Only <reasoning> and <summary> tokens contribute to the loss. The <input> paragraph is masked (label = -100) so the model learns to generate the graph and summary, not memorise the input.

Hyperparameters

Training curves

Train and validation loss stayed within ~0.03 throughout — no overfitting.

---

Limitations

• Trained on Wikipedia-style encyclopaedic paragraphs; may produce lower-quality graphs on conversational or highly technical text
• 360M parameters — graph structure may be incomplete or inconsistent on long or complex inputs
• Max context 1024 tokens; paragraphs longer than ~700 words will be truncated