test

newtest.py

@@ -40,7 +40,7 @@ trainer = SetFitTrainer(
     loss_class=CosineSimilarityLoss,
     metric="accuracy",
     batch_size=64,
-    num_iterations=1,  # 20, # The number of text pairs to generate for contrastive learning
+    num_iterations= 20, # The number of text pairs to generate for contrastive learning
     num_epochs=1,  # The number of epochs to use for constrastive learning
 )
 
@@ -53,10 +53,12 @@ print(f"train dataset: {len(train_dataset)} samples")
 print(f"accuracy: {metrics['accuracy']}")
 
 # Push model to the Hub
-trainer.push_to_hub("MyFirstModel")
+trainer.model.save_pretrained("my_first_test")
 
 # Download from Hub and run inference
-model = SetFitModel.from_pretrained("VinceItsMe/MyFirstModel")
+model = SetFitModel.from_pretrained("my_first_test")
 # Run inference
-preds = model(["i loved the spiderman movie!", "pineapple on pizza is the worst 🤮"])
+preds = model(["i loved France!", "pineapple on pizza is the worst when watching football"])
+label = {'0': 'World','1': 'Sports', '2': 'Business', '3': 'Sci/Tech'}
+output = [label[str(tt.item())] for tt in preds]
 q = 1