discojs/training/disco: add train helpers

epfml · Jul 2, 2024 · 4a86199 · 4a86199
1 parent 425c83c
commit 4a86199
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/docs/examples/training.ts b/docs/examples/training.ts
@@ -14,9 +14,7 @@ async function runUser (url: URL, task: Task, dataset: data.DataSplit): Promise<
   const disco = new Disco(task, { url, scheme: 'federated' })
 
   // Run training on the dataset
-  for await (const round of disco.fit(dataset))
-    for await (const epoch of round)
-      for await (const _ of epoch);
+  await disco.trainFully(dataset);
 
   // Disconnect from the remote server
   await disco.close()

diff --git a/docs/examples/wikitext.ts b/docs/examples/wikitext.ts
@@ -29,9 +29,7 @@ async function main(): Promise<void> {
     const aggregator = new aggregators.MeanAggregator()
     const client = new clients.federated.FederatedClient(url, task, aggregator)
     const disco = new Disco(task, { scheme: 'federated', client, aggregator })
-    for await (const round of disco.fit(dataset))
-      for await (const epoch of round)
-        for await (const _ of epoch);
+    await disco.trainFully(dataset);
 
     // Get the model and complete the prompt
     if (aggregator.model === undefined) {