def sample_rule(rules, nonterminal_coef=1):
  """Sample a rule. Rules with nonterminals are scaled with nonterminal_coef."""
  if not rules:
    raise ValueError

  if len(rules) == 1:
    return rules[0]

  weights = [
      1 if cfg_rule.get_arity(rule) == 0 else nonterminal_coef for rule in rules
  ]
  sampled_idx = random.choices(range(len(rules)), weights=weights, k=1)[0]
  sampled_rule = rules[sampled_idx]
  return sampled_rule
Sample a rule. Rules with nonterminals are scaled with nonterminal_coef.
language/compgen/csl/cky/cfg_sampler.py