huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Avoid re-encoding the input in the T5 example. (#875)

Browse Source
This commit is contained in:
Laurent Mazare
2023-09-17 11:25:54 +02:00
committed by GitHub
parent eeb54716dd
commit 7f65af1f0d
2 changed files with 17 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
candle-transformers/src/models/t5.rs
View File

@ -636,11 +636,18 @@ impl T5ForConditionalGeneration {
})
}
pub fn forward(&mut self, input_ids: &Tensor, decoder_input_ids: &Tensor) -> Result<Tensor> {
let encoder_output = self.encoder.forward(input_ids, None)?;
pub fn encode(&mut self, input_ids: &Tensor) -> Result<Tensor> {
self.encoder.forward(input_ids, None)
}
pub fn decode(
&mut self,
decoder_input_ids: &Tensor,
encoder_output: &Tensor,
) -> Result<Tensor> {
let decoder_output = self
.decoder
.forward(decoder_input_ids, Some(&encoder_output))?;
.forward(decoder_input_ids, Some(encoder_output))?;
let sequence_output = decoder_output
.narrow(1, decoder_output.dim(1)? - 1, 1)?
.squeeze(1)?;
@ -651,6 +658,11 @@ impl T5ForConditionalGeneration {
Ok(output)
}
pub fn forward(&mut self, input_ids: &Tensor, decoder_input_ids: &Tensor) -> Result<Tensor> {
let encoder_output = self.encode(input_ids)?;
self.decode(decoder_input_ids, &encoder_output)
}
pub fn device(&self) -> &Device {
&self.device
}