darpanaswal commited on
Commit
4d1a2d7
·
verified ·
1 Parent(s): d807d27

Update cross_encoder_reranking_train.py

Browse files
Files changed (1) hide show
  1. cross_encoder_reranking_train.py +4 -4
cross_encoder_reranking_train.py CHANGED
@@ -1,19 +1,19 @@
1
  import os
2
  import json
 
3
  import argparse
4
  import numpy as np
5
- import torch
6
- import torch.nn.functional as F
7
  from tqdm import tqdm
8
  from torch import Tensor
9
- from sentence_transformers import SentenceTransformer
10
  from transformers import AutoTokenizer, AutoModel
11
  from sklearn.cluster import AgglomerativeClustering
 
12
  from sklearn.metrics.pairwise import cosine_similarity
13
 
14
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
15
  # Load embedder once
16
- # embedder = SentenceTransformer("all-MiniLM-L6-v2").to(device)
17
  embedder = SentenceTransformer("AI-Growth-Lab/PatentSBERTa").to(device)
18
 
19
 
 
1
  import os
2
  import json
3
+ import torch
4
  import argparse
5
  import numpy as np
 
 
6
  from tqdm import tqdm
7
  from torch import Tensor
8
+ import torch.nn.functional as F
9
  from transformers import AutoTokenizer, AutoModel
10
  from sklearn.cluster import AgglomerativeClustering
11
+ from sentence_transformers import SentenceTransformer
12
  from sklearn.metrics.pairwise import cosine_similarity
13
 
14
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
15
  # Load embedder once
16
+ # embedder = SentenceTransformer("all-MiniLM-L12-v2").to(device)
17
  embedder = SentenceTransformer("AI-Growth-Lab/PatentSBERTa").to(device)
18
 
19