Update README.md

README.md

@@ -4,30 +4,27 @@ This model was trained with the intent to quickyl classify whether or not an ima
 # Model Details
 ## How to Get Started with the Model
 ```python
-import torch
-import transformers
+from PIL import Image
+import requests
 
-model = transformers.AutoModelForCausalLM.from_pretrained(
-    "miguelcarv/phi-1_5-slimorca",
-    trust_remote_code=True
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+model = AutoModelForImageClassification.from_pretrained(
+    "miguelcarv/resnet-50-text-detector",
 )
-tokenizer = transformers.AutoTokenizer.from_pretrained("microsoft/phi-1_5")
 
+processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50", do_resize=False)
 
-SYSTEM_PROMPT = "You are an AI assistant. You will be given a task. You must generate a detailed and long answer."
-input_text = f"""{SYSTEM_PROMPT}
+url = "http://images.cocodataset.org/train2017/000000044520.jpg"
+image = Image.open(requests.get(url, stream=True).raw).convert('RGB').resize((256,256))
 
-Instruction: Give me the first 5 prime numbers and explain what prime numbers are.
-Output:"""
+inputs = processor(image, return_tensors="pt").pixel_values
 
-with torch.no_grad():
-    outputs = model.generate(
-        tokenizer(input_text, return_tensors="pt")['input_ids'],
-        max_length=1024,
-        num_beams = 3,
-        eos_token_id = tokenizer.eos_token_id
-    )
-print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+outputs = model(inputs)
+logits_per_image = outputs.logits 
+probs = logits_per_image.softmax(dim=1) 
+print(probs)
+# tensor([[0.1149, 0.8851]])
 ```
 # Training Details
 - Trained for three epochs