small

app.py

@@ -136,13 +136,13 @@ class RepLlamaModel:
             batch_dict = create_batch_dict(self.tokenizer, batch_texts, always_add_eos="last")
             batch_dict = {key: value.cuda() for key, value in batch_dict.items()}
 
-            # with torch.cuda.amp.autocast():
-            with torch.no_grad():
-                outputs = self.model(**batch_dict)
-                embeddings = pool(outputs.last_hidden_state, batch_dict['attention_mask'], 'last')
-                embeddings = F.normalize(embeddings, p=2, dim=-1)
-                logger.info(f"Encoded shape: {embeddings.shape}, Norm of first embedding: {torch.norm(embeddings[0]).item()}")
-                all_embeddings.append(embeddings.cpu().numpy())
+            with torch.cuda.amp.autocast():
+                with torch.no_grad():
+                    outputs = self.model(**batch_dict)
+                    embeddings = pool(outputs.last_hidden_state, batch_dict['attention_mask'], 'last')
+                    embeddings = F.normalize(embeddings, p=2, dim=-1)
+                    logger.info(f"Encoded shape: {embeddings.shape}, Norm of first embedding: {torch.norm(embeddings[0]).item()}")
+                    all_embeddings.append(embeddings.cpu().numpy())
 
         self.model = self.model.cpu()
         return np.concatenate(all_embeddings, axis=0)