temperature setting fix

microsoft · May 9, 2024 · 5291046 · 5291046
1 parent fd8c5ee
commit 5291046
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/rllm/rllm-base/src/engine.rs b/rllm/rllm-base/src/engine.rs
@@ -515,7 +515,7 @@ impl<ME: ModelExec> RllmEngine<ME> {
                                 let seq_idx = b.sample_mask.unwrap();
                                 aici_bias.apply(&mut logits, seq_idx);
                                 if let Some(t) = b.temperature {
-                                    sg.logits_processor.temperature = Some(t);
+                                    sg.logits_processor.set_temperature(t);
                                 }
                             }
                             None => {}

diff --git a/rllm/rllm-base/src/logits.rs b/rllm/rllm-base/src/logits.rs
@@ -24,4 +24,12 @@ impl LogitsProcessor {
             top_p: sampling_params.top_p,
         }
     }
+
+    pub fn set_temperature(&mut self, temperature: f32) {
+        if temperature < SAMPLING_EPS {
+            self.temperature = None;
+        } else {
+            self.temperature = Some(temperature);
+        }
+    }
 }