Skip to content

Commit

Permalink
Correction: Global initialization of the SentenceTransformer model wa…
Browse files Browse the repository at this point in the history
…s redundant
  • Loading branch information
kreeedit committed Jul 26, 2023
1 parent f20f6e6 commit f65af36
Showing 1 changed file with 4 additions and 5 deletions.
9 changes: 4 additions & 5 deletions trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -245,11 +245,6 @@ def compare_texts_minhash(directory, window_size, step_size, ngram_size, similar
return similarities

# SentenceTransformer functions
model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
# Move model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)


def compare_texts_embeddings(directory, window_size, step_size, model_type, similarity_threshold=0.7):
"""
Expand All @@ -270,6 +265,10 @@ def compare_texts_embeddings(directory, window_size, step_size, model_type, simi
# Initialize the SentenceTransformer model
model = SentenceTransformer(model_type)

# Move model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# Count the number of files to process
num_files = len([name for name in os.listdir(directory) if name.endswith(".txt")])

Expand Down

0 comments on commit f65af36

Please sign in to comment.