code: def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]:
      if len(argv) > 1:
          raise app.UsageError("Too many command-line arguments.")
  
      ##
      # Load models, according to the --models flag.
      models = {}
      for model_string in _MODELS.value:
          # Only split on the first ':', because path may be a URL
          # containing 'https://'
          model_name, path = model_string.split(":", 1)
          logging.info("Loading model '%s' from '%s'", model_name, path)
          if model_name.startswith("bert-"):
              models[model_name] = pretrained_lms.BertMLM(
                  path, top_k=_TOP_K.value
              )
          elif model_name.startswith("gpt2") or model_name in [
              "distilgpt2"
          ]:
              models[model_name] = pretrained_lms.GPT2LanguageModel(
                  path, top_k=_TOP_K.value
              )
          else:
              raise ValueError(
                  f"Unsupported model name '{model_name}' from path '{path}'"
              )
  
      datasets = {
          # Single sentences from movie reviews (SST dev set).
          "sst_dev": glue.SST2Data("validation").remap(
              {"sentence": "text"}
          ),
          # Longer passages from movie reviews (IMDB dataset, test split).
          "imdb_train": classification.IMDBData("test"),
          # Empty dataset, if you just want to type sentences into the UI.
          "blank": lm.PlaintextSents(""),
      }
  
      dataset_loaders: lit_app.DatasetLoadersMap = {
          "sst_dev": (
              glue.SST2DataForLM,
              glue.SST2DataForLM.init_spec(),
          ),
          "imdb_train": (
              classification.IMDBData,
              classification.IMDBData.init_spec(),
          ),
          "plain_text_sentences": (
              lm.PlaintextSents,
              lm.PlaintextSents.init_spec(),
          ),
      }
  
      # Guard this with a flag, because TFDS will download and process 1.67 GB
      # of data if you haven't loaded `lm1b` before.
      if _LOAD_BWB.value:
          # A few sentences from the Billion Word Benchmark (Chelba et al. 2013).
          datasets["bwb"] = lm.BillionWordBenchmark(
              "train", max_examples=_MAX_EXAMPLES.value
          )
          dataset_loaders["bwb"] = (
              lm.BillionWordBenchmark,
              lm.BillionWordBenchmark.init_spec(),
          )
  
      for name in datasets:
          datasets[name] = datasets[name].slice[: _MAX_EXAMPLES.value]
          logging.info(
              "Dataset: '%s' with %d examples",
              name,
              len(datasets[name]),
          )
  
      generators = {"word_replacer": word_replacer.WordReplacer()}
  
      lit_demo = dev_server.Server(
          models,
          datasets,
          generators=generators,
          layouts=CUSTOM_LAYOUTS,
          dataset_loaders=dataset_loaders,
          **server_flags.get_flags(),
      )
      return lit_demo.serve()
  
file_path: lit/lit_nlp/examples/lm_demo.py
imports: from collections.abc import Sequence
  import sys
  from typing import Optional
  from absl import app
  from absl import flags
  from absl import logging
  from lit_nlp import app
  from lit_nlp import dev_server
  from lit_nlp import server_flags
  from lit_nlp.api import layout
  from lit_nlp.components import word_replacer
  from lit_nlp.examples.datasets import classification
  from lit_nlp.examples.datasets import glue
  from lit_nlp.examples.datasets import lm
  from lit_nlp.examples.models import pretrained_lms
label_desc: Use image node_id 2 for calling a global function with example usage: main(argv) and returns: lit_demo