Update 2.py

2.py

@@ -76,7 +76,7 @@ def create_model_from_folder(folder_path):
                 s[os.path.basename(r).replace('.', '_')].extend(parse_xml(os.path.join(r, file)))
     return DM(dict(s))
 
-def create_embeddings_and_sentences(folder_path, model_name="sentence-transformers/all-MiniLM-L6-v2"):
+def create_embeddings_and_sentences(folder_path, model_name="pile-of-law/legalbert-large-1.7M-1"):
     t, m, embeddings, ds = AutoTokenizer.from_pretrained(model_name), AutoModel.from_pretrained(model_name), [], []
     for r, d, f in os.walk(folder_path):
         for file in f:
@@ -91,7 +91,7 @@ def create_embeddings_and_sentences(folder_path, model_name="sentence-transforme
                         ds.append(text)
     return np.vstack(embeddings), ds
 
-def query_vector_similarity(query, embeddings, ds, model_name="sentence-transformers/all-MiniLM-L6-v2"):
+def query_vector_similarity(query, embeddings, ds, model_name="pile-of-law/legalbert-large-1.7M-2"):
     t, m = AutoTokenizer.from_pretrained(model_name), AutoModel.from_pretrained(model_name)
     i = t(query, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():