Update README.md

README.md

@@ -252,13 +252,21 @@ The following hyperparameters were used during training:
 
 ### Example of usage
 
+Simple demo for Google Colab
+
 ```python
-from datasets import load_dataset
+!pip install datasets transformers[torch] accelerate -U
+!git clone https://github.com/Andron00e/CLIPForImageClassification
+%cd CLIPForImageClassification/clip_for_classification
+
+import torch
 from transformers import TrainingArguments
+from datasets import load_dataset, load_metric
 from transformers import CLIPProcessor, AutoModelForImageClassification
+from modeling_clipforimageclassification import CLIPForImageClassification
 
 processor = CLIPProcessor.from_pretrained("Andron00e/CLIPForImageClassification-v1")
-model = AutoModelForImageClassification.from_pretrained("Andron00e/CLIPForImageClassification-v1")
+model = CLIPForImageClassification.from_pretrained("Andron00e/CLIPForImageClassification-v1", 10)
 
 dataset = load_dataset("Andron00e/CIFAR10-custom")
 dataset = dataset["train"].train_test_split(test_size=0.2)
@@ -271,6 +279,8 @@ dataset = DatasetDict({
     "test": val_test["test"],
 })
 
+classes = {0: "airplane", 1: "automobile", 2: "bird", 3: "cat", 4: "deer", 5: "dog", 6: "frog", 7: "horse", 8: "ship", 9: "truck"}
+
 def transform(example_batch):
     inputs = processor(text=[classes[x] for x in example_batch['labels']], images=[x for x in example_batch['image']], padding=True, return_tensors='pt')
     inputs['labels'] = example_batch['labels']
@@ -284,6 +294,11 @@ def collate_fn(batch):
         'labels': torch.tensor([x['labels'] for x in batch])
     }
 
+metric = load_metric("accuracy")
+
+def compute_metrics(p):
+    return metric.compute(predictions=np.argmax(p.predictions, axis=1), references=p.label_ids)
+
 training_args = TrainingArguments(
   output_dir="./outputs",
   per_device_train_batch_size=16,
@@ -322,4 +337,7 @@ trainer.save_state()
 metrics = trainer.evaluate(processed_dataset['test'])
 trainer.log_metrics("eval", metrics)
 trainer.save_metrics("eval", metrics)
+
+%cd ..
+%cd ..
 ```