Update README.md

README.md

@@ -17,6 +17,36 @@ This modelcard aims to be a base template for new models. It has been generated
 
 <!-- Provide a longer summary of what this model is. -->
 
+```
+class WhisperCTC(nn.Module):
+    def __init__(
+        self,
+        encoder_id: str = "tuanio/whisper-encoder.tiny.en",
+        dropout: float = 0.1,
+        vocab_size: int = 47,
+    ):
+        super().__init__()
+        self.encoder = WhisperEncoder.from_pretrained(encoder_id)
+        print("Freezing Whisper Encoder...")
+        self.encoder._freeze_parameters()
+        print("Freezed!")
+        self.lm_head = nn.Sequential(
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Linear(self.encoder.config.d_model, vocab_size),
+        )
+        nn.init.kaiming_uniform_(
+            self.lm_head[-1].weight, mode="fan_in", nonlinearity="relu"
+        )
+
+    def forward(self, feat: Tensor, attn_mask: Tensor):
+        enc = self.encoder(
+            input_features=feat, attention_mask=attn_mask
+        ).last_hidden_state
+        logits = self.lm_head(enc)
+        log_probs = nn.functional.log_softmax(logits, dim=-1)
+        return log_probs
+```
 
 
 - **Developed by:** [More Information Needed]