Refactor temp

armbues · Jun 28, 2024 · 89549ac · 89549ac
1 parent 5ccfe17
commit 89549ac
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/sillm/core/llm.py b/sillm/core/llm.py
@@ -406,10 +406,10 @@ def generate(model,
  buffer, output = [], ""
 
  def sample(logits):
- if temperature > 0:
- return mx.random.categorical(logits * (1 / temperature))
- else:
+ if temperature == 0:
  return mx.argmax(logits, axis=-1)
+ else:
+ return mx.random.categorical(logits * (1 / temperature))
  # TODO add top-p sampling
 
  def apply_repetition_penalty(logits):