Update README.md

README.md

@@ -56,6 +56,7 @@ Start by installing the [library](https://huggingface.co/docs/transformers): `pi
 
 ```python
 import torch
+from torch.nn.functional import relu, normalize
 from transformers import AutoTokenizer, AutoModel
 
 queries = ["Ceci est un exemple de requête.", "Voici un second exemple."]
@@ -71,11 +72,11 @@ with torch.no_grad():
     q_output = model(**q_input)
     p_output = model(**p_input)
 
-q_activations = torch.amax(torch.log1p(input=self.relu(q_output.logits * q_input['attention_mask'].unsqueeze(-1))), dim=1)
-p_activations = torch.amax(torch.log1p(input=self.relu(p_output.logits * p_input['attention_mask'].unsqueeze(-1))), dim=1)
+q_activations = torch.amax(torch.log1p(relu(q_output.logits * q_input['attention_mask'].unsqueeze(-1))), dim=1)
+p_activations = torch.amax(torch.log1p(relu(p_output.logits * p_input['attention_mask'].unsqueeze(-1))), dim=1)
 
-q_activations = torch.nn.functional.normalize(q_activations, p=2, dim=1)
-p_activations = torch.nn.functional.normalize(p_activations, p=2, dim=1)
+q_activations = normalize(q_activations, p=2, dim=1)
+p_activations = normalize(p_activations, p=2, dim=1)
 
 similarity = q_embeddings @ p_embeddings.T
 print(similarity)