Predict EOS probability. Input: (batch, hidden_size) LLM hidden state. Output: (batch, 1) probability (after sigmoid).
(&self, x: &Tensor)
| 34 | /// Input: (batch, hidden_size) LLM hidden state. |
| 35 | /// Output: (batch, 1) probability (after sigmoid). |
| 36 | pub fn forward(&self, x: &Tensor) -> Result<Tensor> { |
| 37 | let h = self.backend.linear_forward(x, &self.fc1_weight, self.fc1_bias.as_ref())?; |
| 38 | let h = self.backend.silu(&h)?; |
| 39 | let h = self.backend.linear_forward(&h, &self.fc2_weight, self.fc2_bias.as_ref())?; |
| 40 | self.backend.sigmoid(&h) |
| 41 | } |
| 42 | |
| 43 | /// Check if generation should stop (probability > threshold). |
| 44 | pub fn should_stop(&self, x: &Tensor, threshold: f32) -> Result<bool> { |