effect-gpt 0.1.0

package/README.md

@@ -0,0 +1,50 @@
+# Effect GPT
+
+A transformer-based LLM built from scratch with Effect. Inspired by [RustGPT](https://github.com/tekaratzas/RustGPT).
+
+## What This Is
+
+A complete LLM implementation including:
+- **Tokenization** — BPE-style text preprocessing
+- **Transformer Architecture** — embeddings, multi-head attention, feed-forward layers, layer norm
+- **Training** — cross-entropy loss, backpropagation, Adam optimizer with gradient clipping
+- **Inference** — greedy decoding for text generation
+
+## Quick Start
+
+```bash
+bun install
+bun run dev      # train + generate
+bun test         # run test suite
+```
+
+## Why Effect?
+
+This project leverages Effect's robust ecosystem to bring systems-programming discipline to TypeScript:
+
+- **Service-Based Architecture** — Uses `Context` and `Layer` to keep the core model pure and make testing deterministic by swapping implementations (e.g., swapping a terminal logger for a test capture).
+- **Type-Safe Errors** — Implements `Data.TaggedError` for a precise, union-based error system, ensuring all failure cases (shape mismatches, IO errors) are handled explicitly.
+- **Lazy Streaming** — Utilizes `Stream` for high-performance, backpressured data loading and batching during training.
+- **Resource Management** — Uses `Scope` to guarantee that file handles and fibers are always cleaned up correctly.
+- **Declarative Concurrency** — Leverages Effect's runtime to manage parallel preprocessing and training loops without the complexity of manual orchestration.
+
+## Deep Dive
+
+Curious about how this was built? Check out our blog post: [Building a transformer-based LLM with Effect](https://hackmd.io/jW2kapAMSkWpG_PW1o3hAA).
+
+## Project Structure
+
+```
+src/
+├── tensor/      # Tensor2D, matmul, softmax, layer norm
+├── model/       # embeddings, attention, transformer blocks, forward/backward
+├── training/    # loss, gradients, Adam optimizer, training loop
+├── tokenize/    # text → tokens
+├── vocab/       # vocabulary management
+├── services/    # Effect services (Random, Logger, Metrics)
+└── cli/         # command-line interface
+tests/
+├── ts/                  # TypeScript test suite
+│   ├── support/         # test helpers, seeded RNG, 
+│   └── ...              # model, tensor, and training unit/integration tests
+```

package/data/chat_training_data.json

@@ -0,0 +1,55 @@
+[
+    "User: What causes rain? Assistant: Rain is caused by water vapor in clouds condensing into droplets that become too heavy to remain airborne </s>",
+    "User: How do mountains form? Assistant: Mountains are formed through tectonic forces or volcanism over long geological time periods </s>",
+    "User: What is the Amazon rainforest? Assistant: The Amazon rainforest is one of the most biodiverse places on earth, home to countless species </s>",
+    "User: At what temperature does water boil? Assistant: Water boils at 100 degrees celsius at standard atmospheric pressure </s>",
+    "User: How long does it take the moon to orbit Earth? Assistant: The moon orbits the earth approximately every 27.3 days </s>",
+    "User: What is photosynthesis? Assistant: Photosynthesis is the process by which green plants use sunlight to synthesize food from carbon dioxide </s>",
+    "User: How does gravity work? Assistant: Gravity is a force that attracts two bodies toward each other based on their mass </s>",
+    "User: How many neurons are in the human brain? Assistant: The human brain contains about 86 billion neurons that transmit information throughout the body </s>",
+    "User: What is electricity? Assistant: Electricity is the flow of electrons through a conductor, often used to power devices </s>",
+    "User: What is climate change? Assistant: Climate change refers to long-term shifts in temperatures and weather patterns on Earth </s>",
+    "User: How long do oak trees live? Assistant: Oak trees can live for hundreds of years and produce acorns as their fruit </s>",
+    "User: What happened to Pluto? Assistant: Pluto was reclassified from a planet to a dwarf planet in 2006 by astronomers </s>",
+    "User: How is glass made? Assistant: Glass is made by heating sand, soda ash, and limestone to very high temperatures until they melt </s>",
+    "User: What do volcanoes do? Assistant: Volcanoes can erupt with lava, ash, and gases, altering landscapes and ecosystems </s>",
+    "User: Why was the Great Wall built? Assistant: The great wall of china was built to protect ancient china from invasions </s>",
+    "User: What are penguins like? Assistant: Penguins are flightless birds that are well adapted to life in cold environments </s>",
+    "User: How much rain do deserts get? Assistant: Deserts receive less than 250 millimeters of precipitation each year </s>",
+    "User: What is Jupiter like? Assistant: Jupiter is the largest planet in our solar system and has dozens of moons </s>",
+    "User: How fast does light travel? Assistant: Light travels at approximately 299,792 kilometers per second in a vacuum </s>",
+    "User: What is gold used for? Assistant: Gold is a dense, soft metal often used in jewelry and electronics due to its conductivity </s>",
+    "User: What covers most of Earth? Assistant: Most of the earth's surface is covered by water, primarily in oceans </s>",
+    "User: How do bicycles work? Assistant: Bicycles are an efficient mode of transport that convert human energy into motion through gears </s>",
+    "User: How is chocolate made? Assistant: Chocolate is made from roasted and ground cacao seeds, often sweetened and flavored </s>",
+    "User: What is the internet? Assistant: The internet is a global network that allows for digital communication and information sharing </s>",
+    "User: How do wind turbines work? Assistant: Wind energy is harnessed using turbines that convert wind motion into electricity </s>",
+    "User: What are cats like? Assistant: Cats are domesticated mammals known for their independence and hunting instincts </s>",
+    "User: How do languages change? Assistant: Languages evolve over time through cultural, social, and technological influences </s>",
+    "User: What did the printing press do? Assistant: The printing press revolutionized the spread of information in the 15th century </s>",
+    "User: What is sound? Assistant: Sound is a vibration that travels through air, water, or solid materials to reach our ears </s>",
+    "User: Why is carbon important? Assistant: Carbon is an essential element in organic chemistry, forming the basis of all life on Earth </s>",
+    "User: What was the Library of Alexandria? Assistant: The library of alexandria was one of the most significant libraries of the ancient world </s>",
+    "User: What do honeybees do? Assistant: Honeybees play a vital role in pollination, which supports ecosystems and agriculture </s>",
+    "User: Are electric cars better? Assistant: Electric vehicles produce less air pollution than traditional gasoline-powered cars </s>",
+    "User: How do you make bread? Assistant: Bread is typically made from flour, water, yeast, and salt through a baking process </s>",
+    "User: What is the Sahara desert? Assistant: The sahara desert is the largest hot desert in the world, spanning multiple countries </s>",
+    "User: What are renewable resources? Assistant: Renewable resources replenish naturally and include sunlight, wind, and water </s>",
+    "User: What are eclipses? Assistant: Eclipses occur when one celestial body moves into the shadow of another </s>",
+    "User: How are language models trained? Assistant: Language models are trained using vast amounts of text to learn patterns in language </s>",
+    "User: How do compasses work? Assistant: Compasses work by aligning a magnetic needle with the earth's magnetic field </s>",
+    "User: What do vaccines do? Assistant: Vaccines help the immune system recognize and fight off specific pathogens </s>",
+    "User: Hello! Assistant: Hello! How can I help you today? </s>",
+    "User: Hi! Assistant: Hi there! What can I do for you? </s>",
+    "User: Hey! Assistant: Hey! How's it going? What would you like to know? </s>",
+    "User: Good morning! Assistant: Good morning! Hope you're having a great day. How can I assist you? </s>",
+    "User: Good afternoon! Assistant: Good afternoon! What can I help you with today? </s>",
+    "User: Good evening! Assistant: Good evening! How may I help you tonight? </s>",
+    "User: How are you? Assistant: I'm doing well, thank you for asking! How can I help you today? </s>",
+    "User: What's up? Assistant: Not much, just here to help! What would you like to know? </s>",
+    "User: Greetings! Assistant: Greetings! It's nice to meet you. How can I be of assistance? </s>",
+    "User: Thank you. Assistant: You're welcome! I'm happy to help </s>",
+    "User: Can you explain that? Assistant: Of course! I'd be happy to explain in more detail </s>",
+    "User: I don't understand. Assistant: No problem! Let me try to explain it differently </s>",
+    "User: What do you think? Assistant: Based on the information available, I think it's important to consider multiple perspectives </s>"
+]

package/data/pretraining_data.json

@@ -0,0 +1,27 @@
+[
+    "The sun rises in the east and sets in the west </s>",
+    "Water flows downhill due to gravity </s>",
+    "Birds fly through the air using their wings </s>",
+    "Fish swim in rivers, lakes, and oceans </s>",
+    "Trees grow tall and produce leaves </s>",
+    "Rain falls from clouds in the sky </s>",
+    "Fire is hot and produces light </s>",
+    "Ice is frozen water that melts when heated </s>",
+    "Mountains are tall and rocky formations </s>",
+    "The moon orbits around planet Earth </s>",
+    "Flowers bloom in spring and summer </s>",
+    "Snow is cold and white </s>",
+    "Wind moves air from place to place </s>",
+    "Rivers flow into larger bodies of water </s>",
+    "Sand is found on beaches and in deserts </s>",
+    "Grass grows in fields and yards </s>",
+    "Rocks are hard and can be different colors </s>",
+    "Stars shine bright in the night sky </s>",
+    "Waves move across the surface of water </s>",
+    "Clouds form when water vapor rises </s>",
+    "Lightning is bright and makes thunder </s>",
+    "Storms bring rain and strong winds </s>",
+    "Seasons change throughout the year </s>",
+    "Animals eat food to survive </s>",
+    "Plants need sunlight and water to grow </s>"
+]

package/package.json

@@ -0,0 +1,25 @@
+{
+  "name": "effect-gpt",
+  "version": "0.1.0",
+  "type": "module",
+  "main": "dist/index.js",
+  "scripts": {
+    "build": "tsc",
+    "check": "tsc --noEmit",
+    "dev": "bun src/cli/main.ts",
+    "test": "bun test",
+    "test:watch": "bun test --watch",
+    "test:unit": "bun test tests/ts"
+  },
+  "dependencies": {
+    "effect": "^3.16.0",
+    "@effect/platform": "^0.87.0",
+    "@effect/platform-bun": "^0.60.22",
+    "@effect/experimental": "^0.44.21",
+    "@effect/sql": "^0.33.21"
+  },
+  "devDependencies": {
+    "@types/bun": "latest",
+    "typescript": "^5.6.0"
+  }
+}

package/src/cli/errors.ts

@@ -0,0 +1,51 @@
+import * as Cause from "effect/Cause"
+import type { TrainingError } from "../errors"
+
+const formatUnknown = (err: unknown): string => {
+  if (Cause.isCause(err)) return Cause.pretty(err)
+  if (err instanceof Error) return err.stack ?? `${err.name}: ${err.message}`
+  return String(err)
+}
+
+export const formatTrainingError = (err: TrainingError | unknown): string => {
+  if (err && typeof err === "object" && "_tag" in err) {
+    const tagged = err as { _tag: string }
+    switch (tagged._tag) {
+      case "TrainingDatasetError": {
+        const e = err as TrainingError & { cause: { path?: string; error?: unknown; message?: string; _tag?: string } }
+        const location = e.cause.path ? ` (${e.cause.path})` : ""
+        const reason = e.cause._tag ?? "dataset"
+        const detail =
+          e.cause.error !== undefined
+            ? typeof e.cause.error === "object" && e.cause.error !== null && "message" in e.cause.error
+              ? (e.cause.error as any).message
+              : String(e.cause.error)
+            : e.cause.message ?? ""
+        return `Dataset error${location}: ${reason}${detail ? ` - ${detail}` : ""}`
+      }
+      case "TrainingShapeError": {
+        const e = err as any
+        return `Shape error: ${e.cause?.message ?? formatUnknown(e.cause)}`
+      }
+      case "TrainingTokenizerError": {
+        const e = err as any
+        return `Tokenizer error: ${e.message}`
+      }
+      case "TrainingOptimizerError": {
+        const e = err as any
+        return `Optimizer error: ${e.message}`
+      }
+      case "TrainingConfigError": {
+        const e = err as any
+        return `Configuration error: ${e.message}`
+      }
+      case "TrainingUnknownError": {
+        const e = err as any
+        return `Unexpected training error: ${formatUnknown(e.cause)}`
+      }
+      default:
+        return `Unexpected error (${tagged._tag}): ${formatUnknown(err)}`
+    }
+  }
+  return `Unexpected error: ${formatUnknown(err)}`
+}

package/src/cli/main.ts

@@ -0,0 +1,163 @@
+import * as Effect from "effect/Effect"
+import * as Layer from "effect/Layer"
+import * as HashSet from "effect/HashSet"
+import { Terminal } from "@effect/platform"
+import { BunFileSystem, BunRuntime, BunTerminal } from "@effect/platform-bun"
+import { Dataset } from "../data/Dataset"
+import { Vocab } from "../vocab/Vocab"
+import { LLM } from "../model/LLM"
+import { Embeddings } from "../model/Embeddings"
+import { TransformerBlock } from "../model/TransformerBlock"
+import { OutputProjection } from "../model/OutputProjection"
+import {
+  trainStream,
+  makeLLMLayer,
+  makeTrainingConfigLayer,
+  makePreprocessSettingsLayer
+} from "../training/train"
+import { MAX_SEQ_LEN, EMBEDDING_DIM, HIDDEN_DIM } from "../config"
+import { PrettyLoggerLive, info, error as logError } from "../services/Logger"
+import { InMemoryMetricsLive, snapshot } from "../services/Metrics"
+import { SeedLayer, useSeedRng } from "../services/SeedLayer"
+import type { Rng } from "../tensor/random"
+import type { TrainingError } from "../errors"
+import { formatTrainingError } from "./errors"
+
+const PRETRAIN_EPOCHS = 100
+const PRETRAIN_LR = 0.0005
+const FINETUNE_EPOCHS = 100
+const FINETUNE_LR = 0.0001
+
+const readLine = (prompt: string) =>
+  Effect.gen(function* () {
+    const terminal = yield* Terminal.Terminal
+    yield* terminal.display(prompt)
+    return yield* terminal.readLine
+  })
+
+const repl = (llm: LLM) =>
+  Effect.scoped(
+    Effect.gen(function* () {
+      const terminal = yield* Terminal.Terminal
+      yield* terminal.display("\n--- Interactive Mode ---\n")
+      yield* terminal.display("Type a prompt and press Enter to generate text.\n")
+      yield* terminal.display("Type 'exit' to quit.\n")
+
+      while (true) {
+        const input = yield* readLine("\nEnter prompt: ")
+        const trimmed = input.trim()
+
+        if (trimmed.toLowerCase() === "exit") {
+          yield* terminal.display("Exiting interactive mode.\n")
+          break
+        }
+
+        const formattedInput = `User: ${trimmed}`
+        const prediction = yield* llm.predict(formattedInput)
+        yield* terminal.display(`Model output: ${prediction}\n`)
+      }
+    })
+  )
+
+const parseSeedArg = (argv: string[]): number | undefined => {
+  const seedIndex = argv.findIndex((arg) => arg === "--seed")
+  if (seedIndex >= 0 && seedIndex < argv.length - 1) {
+    const asNum = Number(argv[seedIndex + 1])
+    return Number.isFinite(asNum) ? asNum : undefined
+  }
+  return undefined
+}
+
+const main = Effect.scoped(
+  Effect.gen(function* () {
+    const terminal = yield* Terminal.Terminal
+    const rng: Rng = yield* useSeedRng()
+
+    const dataset = Dataset.load({
+      pretrainingPath: "data/pretraining_data.json",
+      chatPath: "data/chat_training_data.json",
+      format: "json"
+    })
+
+    const vocabSet1 = yield* Vocab.processStreamForVocab(dataset.pretrainingStream())
+    const vocabSet2 = yield* Vocab.processStreamForVocab(dataset.chatStream())
+
+    const combinedSet = HashSet.union(vocabSet1, vocabSet2)
+    const vocabWords = Array.from(HashSet.values(combinedSet)).sort()
+    const vocab = Vocab.make(vocabWords)
+
+    const vocabSize = vocab.words.length
+    const network = [
+      new Embeddings(vocabSize, EMBEDDING_DIM, MAX_SEQ_LEN, rng),
+      new TransformerBlock(EMBEDDING_DIM, HIDDEN_DIM, rng),
+      new TransformerBlock(EMBEDDING_DIM, HIDDEN_DIM, rng),
+      new TransformerBlock(EMBEDDING_DIM, HIDDEN_DIM, rng),
+      new OutputProjection(EMBEDDING_DIM, vocabSize, rng)
+    ]
+    const llm = new LLM(vocab, network)
+
+    yield* terminal.display("\n=== MODEL INFORMATION ===\n")
+    yield* terminal.display(`Network architecture: ${llm.networkDescription()}\n`)
+    yield* terminal.display(
+      `Model configuration -> max_seq_len: ${MAX_SEQ_LEN}, embedding_dim: ${EMBEDDING_DIM}, hidden_dim: ${HIDDEN_DIM}\n`
+    )
+    yield* terminal.display(`Total parameters: ${llm.totalParameters()}\n`)
+
+    const testInput = "User: How do mountains form?"
+
+    yield* terminal.display("\n=== BEFORE TRAINING ===\n")
+    yield* terminal.display(`Input: ${testInput}\n`)
+    const beforeOutput = yield* llm.predict(testInput)
+    yield* terminal.display(`Output: ${beforeOutput}\n`)
+
+    const llmLayer = makeLLMLayer(llm)
+    const preprocessLayer = makePreprocessSettingsLayer({ concurrency: "unbounded", batchSize: 1 })
+
+    yield* info("\n=== PRE-TRAINING MODEL ===")
+    yield* info(`Pre-training for ${PRETRAIN_EPOCHS} epochs with learning rate ${PRETRAIN_LR}`)
+    yield* trainStream(dataset.pretrainingStream).pipe(
+      Effect.provide(llmLayer),
+      Effect.provide(makeTrainingConfigLayer({ epochs: PRETRAIN_EPOCHS, learningRate: PRETRAIN_LR })),
+      Effect.provide(preprocessLayer)
+    )
+
+    yield* info("\n=== INSTRUCTION TUNING ===")
+    yield* info(`Instruction tuning for ${FINETUNE_EPOCHS} epochs with learning rate ${FINETUNE_LR}`)
+    yield* trainStream(dataset.chatStream).pipe(
+      Effect.provide(llmLayer),
+      Effect.provide(makeTrainingConfigLayer({ epochs: FINETUNE_EPOCHS, learningRate: FINETUNE_LR })),
+      Effect.provide(preprocessLayer)
+    )
+
+    const metrics = yield* snapshot()
+    yield* info("Training complete", {
+      epochsCompleted: metrics.counters.find((c) => c.name === "epochs_completed")?.value,
+      totalExamples: metrics.counters.find((c) => c.name === "examples_processed")?.value,
+      finalLoss: metrics.gauges.find((g) => g.name === "epoch_loss")?.value
+    })
+
+    yield* terminal.display("\n=== AFTER TRAINING ===\n")
+    yield* terminal.display(`Input: ${testInput}\n`)
+    const afterOutput = yield* llm.predict(testInput)
+    yield* terminal.display(`Output: ${afterOutput}\n`)
+    yield* terminal.display("======================\n")
+
+    yield* repl(llm)
+  })
+)
+
+const LoggerLayer = PrettyLoggerLive("info")
+
+const seedValue = parseSeedArg(process.argv)
+const SeedLayerLive = SeedLayer(seedValue)
+
+const AppLayer = Layer.mergeAll(BunFileSystem.layer, BunTerminal.layer, LoggerLayer, InMemoryMetricsLive, SeedLayerLive)
+
+const program = Effect.scoped(
+  main.pipe(
+    Effect.provide(AppLayer),
+    Effect.catchAll((err) => logError(formatTrainingError(err as TrainingError)).pipe(Effect.provide(AppLayer)))
+  )
+)
+
+BunRuntime.runMain(program)

package/src/config.ts

@@ -0,0 +1,3 @@
+export const MAX_SEQ_LEN = 80
+export const EMBEDDING_DIM = 128
+export const HIDDEN_DIM = 256

package/src/data/Dataset.ts

@@ -0,0 +1,168 @@
+import * as Data from "effect/Data"
+import * as Effect from "effect/Effect"
+import * as Stream from "effect/Stream"
+import * as Option from "effect/Option"
+import * as Schema from "effect/Schema"
+import { FileSystem } from "@effect/platform"
+
+export type DatasetFormat = "json" | "csv"
+
+export interface DatasetLoadOptions {
+  readonly pretrainingPath: string
+  readonly chatPath: string
+  readonly format: DatasetFormat
+}
+
+export interface DatasetStreams {
+  /**
+   * Fresh stream producer for pretraining data.
+   * Each invocation re-opens the file, so it is safe to call once per epoch.
+   */
+  readonly pretrainingStream: () => Stream.Stream<string, DatasetLoadError | DatasetParseError, FileSystem.FileSystem>
+  /**
+   * Fresh stream producer for chat training data.
+   * Each invocation re-opens the file, so it is safe to call once per epoch.
+   */
+  readonly chatStream: () => Stream.Stream<string, DatasetLoadError | DatasetParseError, FileSystem.FileSystem>
+}
+
+export interface DatasetArrays {
+  readonly pretrainingData: ReadonlyArray<string>
+  readonly chatTrainingData: ReadonlyArray<string>
+}
+
+export class DatasetLoadError extends Data.TaggedError("DatasetLoadError")<{
+  readonly path: string
+  readonly error: unknown
+}> {}
+
+export class DatasetParseError extends Data.TaggedError("DatasetParseError")<{
+  readonly path: string
+  readonly error: unknown
+}> {}
+
+const TrainingItemSchema = Schema.String
+const decodeTrainingItem = Schema.decodeUnknown(TrainingItemSchema)
+
+const makeFileStream = (path: string) =>
+  Effect.flatMap(FileSystem.FileSystem, (fs) => Effect.succeed(fs.stream(path)))
+
+const linesFromFile = (path: string) =>
+  Stream.splitLines(
+    Stream.decodeText(
+      Stream.mapError(
+        Stream.unwrap(makeFileStream(path)),
+        (error) => new DatasetLoadError({ path, error })
+      )
+    )
+  )
+
+const parseJsonLine = (path: string, rawLine: string): Effect.Effect<Option.Option<string>, DatasetParseError> =>
+  Effect.gen(function* () {
+    const trimmed = rawLine.trim()
+    if (trimmed === "[" || trimmed === "]" || trimmed.length === 0) {
+      return yield* Effect.succeed<Option.Option<string>>(Option.none())
+    }
+
+    const withoutComma = trimmed.endsWith(",") ? trimmed.slice(0, -1) : trimmed
+
+    const parsed: unknown = yield* Effect.try({
+      try: () => JSON.parse(withoutComma),
+      catch: (error) => new DatasetParseError({ path, error })
+    })
+
+    const decoded: string = yield* decodeTrainingItem(parsed).pipe(
+      Effect.mapError((error) => new DatasetParseError({ path, error }))
+    )
+
+    return yield* Effect.succeed<Option.Option<string>>(Option.some(decoded))
+  })
+
+const jsonStream = (path: string) =>
+  Stream.mapEffect(linesFromFile(path), (line) => parseJsonLine(path, line)).pipe(
+    Stream.filterMap((option) => option)
+  )
+
+const splitCsvLine = (line: string): Array<string> => {
+  const fields: Array<string> = []
+  let current = ""
+  let inQuotes = false
+
+  for (let i = 0; i < line.length; i++) {
+    const ch = line[i]!
+    if (ch === '"') {
+      if (inQuotes && line[i + 1] === '"') {
+        // Escaped quote
+        current += '"'
+        i++ // skip next
+      } else {
+        inQuotes = !inQuotes
+      }
+    } else if (ch === "," && !inQuotes) {
+      fields.push(current)
+      current = ""
+    } else {
+      current += ch
+    }
+  }
+
+  if (inQuotes) {
+    throw new Error("Unclosed quote in CSV line")
+  }
+
+  fields.push(current)
+  return fields
+}
+
+const parseCsvLine = (path: string, rawLine: string): Effect.Effect<Option.Option<string>, DatasetParseError> =>
+  Effect.gen(function* () {
+    const trimmed = rawLine.trim()
+    if (trimmed.length === 0) {
+      return yield* Effect.succeed<Option.Option<string>>(Option.none())
+    }
+
+    const joined = yield* Effect.try({
+      try: () => splitCsvLine(rawLine).join(","),
+      catch: (error) => new DatasetParseError({ path, error })
+    })
+
+    return yield* Effect.succeed<Option.Option<string>>(Option.some(joined))
+  })
+
+const csvStream = (path: string) =>
+  Stream.mapEffect(linesFromFile(path), (line) => parseCsvLine(path, line)).pipe(
+    Stream.filterMap((option) => option)
+  )
+
+const streamForFormat = (path: string, format: DatasetFormat) =>
+  format === "json" ? jsonStream(path) : csvStream(path)
+
+const collectAll = (
+  stream: Stream.Stream<string, DatasetLoadError | DatasetParseError, FileSystem.FileSystem>
+) =>
+  Stream.runCollect(stream).pipe(
+    Effect.map((chunk) => Array.from(chunk)),
+    Effect.catchAll((error) => Effect.fail(error))
+  )
+
+export const Dataset = {
+  /**
+   * Returns dataset streams that can be re-opened for each consumption.
+   */
+  load: (options: DatasetLoadOptions): DatasetStreams => ({
+    pretrainingStream: () => streamForFormat(options.pretrainingPath, options.format),
+    chatStream: () => streamForFormat(options.chatPath, options.format)
+  }),
+
+  /**
+   * Convenience helper to materialize both streams into arrays.
+   * Useful for small datasets or legacy call sites.
+   */
+  collect: (
+    streams: DatasetStreams
+  ): Effect.Effect<DatasetArrays, DatasetLoadError | DatasetParseError, FileSystem.FileSystem> =>
+    Effect.all({
+      pretrainingData: collectAll(streams.pretrainingStream()),
+      chatTrainingData: collectAll(streams.chatStream())
+    })
+}

package/src/errors.ts

@@ -0,0 +1,73 @@
+import * as Data from "effect/Data"
+import type { DatasetLoadError, DatasetParseError } from "./data/Dataset"
+import type { ShapeError } from "./tensor/ops"
+
+/**
+ * Narrow, discriminated error channel for the training pipeline.
+ * Keeps original domain errors attached for rich reporting at the CLI boundary.
+ */
+export class TrainingDatasetError extends Data.TaggedError("TrainingDatasetError")<{
+  readonly cause: DatasetLoadError | DatasetParseError
+}> {}
+
+export class TrainingShapeError extends Data.TaggedError("TrainingShapeError")<{
+  readonly cause: ShapeError
+}> {}
+
+export class TrainingTokenizerError extends Data.TaggedError("TrainingTokenizerError")<{
+  readonly message: string
+  readonly cause?: unknown
+}> {}
+
+export class TrainingOptimizerError extends Data.TaggedError("TrainingOptimizerError")<{
+  readonly message: string
+  readonly cause?: unknown
+}> {}
+
+export class TrainingConfigError extends Data.TaggedError("TrainingConfigError")<{
+  readonly message: string
+  readonly cause?: unknown
+}> {}
+
+export class TrainingUnknownError extends Data.TaggedError("TrainingUnknownError")<{
+  readonly cause: unknown
+}> {}
+
+export type TrainingError =
+  | TrainingDatasetError
+  | TrainingShapeError
+  | TrainingTokenizerError
+  | TrainingOptimizerError
+  | TrainingConfigError
+  | TrainingUnknownError
+
+export const TrainingError = {
+  dataset: (cause: DatasetLoadError | DatasetParseError): TrainingDatasetError =>
+    new TrainingDatasetError({ cause }),
+  shape: (cause: ShapeError): TrainingShapeError => new TrainingShapeError({ cause }),
+  tokenizer: (message: string, cause?: unknown): TrainingTokenizerError =>
+    new TrainingTokenizerError({ message, cause }),
+  optimizer: (message: string, cause?: unknown): TrainingOptimizerError =>
+    new TrainingOptimizerError({ message, cause }),
+  config: (message: string, cause?: unknown): TrainingConfigError =>
+    new TrainingConfigError({ message, cause }),
+  unknown: (cause: unknown): TrainingUnknownError => new TrainingUnknownError({ cause }),
+  fromUnknown: (error: unknown): TrainingError => {
+    if (error instanceof TrainingDatasetError) return error
+    if (error instanceof TrainingShapeError) return error
+    if (error instanceof TrainingTokenizerError) return error
+    if (error instanceof TrainingOptimizerError) return error
+    if (error instanceof TrainingConfigError) return error
+    // Upstream domain errors
+    if (error && typeof error === "object") {
+      const candidate = error as { _tag?: string }
+      if (candidate._tag === "DatasetLoadError" || candidate._tag === "DatasetParseError") {
+        return new TrainingDatasetError({ cause: error as DatasetLoadError | DatasetParseError })
+      }
+      if (candidate._tag === "ShapeError") {
+        return new TrainingShapeError({ cause: error as ShapeError })
+      }
+    }
+    return new TrainingUnknownError({ cause: error })
+  }
+}