Update README.md

README.md

@@ -39,18 +39,30 @@ print(embeddings)
 ## Training Code
 ```python
 from datasets import load_dataset, concatenate_datasets
-from sentence_transformers import SentenceTransformer, InputExample, losses, models, util, datasets
+from sentence_transformers import (
+    SentenceTransformer,
+    InputExample,
+    losses,
+    models,
+    util,
+    datasets,
+)
 from torch.utils.data import DataLoader
 from torch import nn
 import random
 
-word_embedding_model = models.Transformer("KBLab/bert-base-swedish-cased-new", max_seq_length=256)
+word_embedding_model = models.Transformer(
+    "KBLab/bert-base-swedish-cased-new", max_seq_length=256
+)
 pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
 dense_model = models.Dense(
-    in_features=pooling_model.get_sentence_embedding_dimension(), out_features=256, activation_function=nn.Tanh()
+    in_features=pooling_model.get_sentence_embedding_dimension(),
+    out_features=256,
+    activation_function=nn.Tanh(),
 )
 model = SentenceTransformer(modules=[word_embedding_model, pooling_model, dense_model])
 
+
 def pair():
     def norm(x):
         x["label"] = x["label"] / m
@@ -58,7 +70,9 @@ def pair():
 
     dd = []
     for sub in ["swepar", "swesim_relatedness", "swesim_similarity"]:
-        ds = concatenate_datasets([d for d in load_dataset("sbx/superlim-2", sub).values()])
+        ds = concatenate_datasets(
+            [d for d in load_dataset("sbx/superlim-2", sub).values()]
+        )
         if "sentence_1" in ds.features:
             ds = ds.rename_column("sentence_1", "d1")
             ds = ds.rename_column("sentence_2", "d2")
@@ -74,10 +88,15 @@ def pair():
         train_examples.append(InputExample(texts=[d["d1"], d["d2"]], label=d["label"]))
     train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=64)
     train_loss = losses.CosineSimilarityLoss(model)
-    model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=10, warmup_steps=100)
+    model.fit(
+        train_objectives=[(train_dataloader, train_loss)], epochs=10, warmup_steps=100
+    )
+
 
 def nli():
-    ds = concatenate_datasets([d for d in load_dataset("sbx/superlim-2", "swenli").values()])
+    ds = concatenate_datasets(
+        [d for d in load_dataset("sbx/superlim-2", "swenli").values()]
+    )
 
     def add_to_samples(sent1, sent2, label):
         if sent1 not in train_data:
@@ -93,17 +112,33 @@ def nli():
     for sent1, others in train_data.items():
         if len(others[0]) > 0 and len(others[1]) > 0:
             train_samples.append(
-                InputExample(texts=[sent1, random.choice(list(others[0])), random.choice(list(others[1]))])
+                InputExample(
+                    texts=[
+                        sent1,
+                        random.choice(list(others[0])),
+                        random.choice(list(others[1])),
+                    ]
+                )
             )
             train_samples.append(
-                InputExample(texts=[random.choice(list(others[0])), sent1, random.choice(list(others[1]))])
+                InputExample(
+                    texts=[
+                        random.choice(list(others[0])),
+                        sent1,
+                        random.choice(list(others[1])),
+                    ]
+                )
             )
     train_dataloader = datasets.NoDuplicatesDataLoader(train_samples, batch_size=64)
     train_loss = losses.MultipleNegativesRankingLoss(model)
-    model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=1, warmup_steps=100)
+    model.fit(
+        train_objectives=[(train_dataloader, train_loss)], epochs=1, warmup_steps=100
+    )
+
 
 pair()
 nli()
 model.save()
 
+
 ```