Upload infer.py with huggingface_hub

infer.py

@@ -46,7 +46,7 @@ def predict_with_automodel(
         return_tensors="pt",
         truncation=True,
         max_length=max_length,
-    )
+    ).to(model.device)
     with torch.no_grad():
         logits = model(**encoded).logits
         probs = torch.softmax(logits, dim=-1).squeeze(0)
@@ -57,9 +57,9 @@ def predict_with_automodel(
 
 
 if __name__ == "__main__":
-    model_dir = "artifacts/models/langid-mmbert-small-8gpu"
+    model_dir = "chiennv/langid-mmbert-small"
     tokenizer = AutoTokenizer.from_pretrained(model_dir)
-    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+    model = AutoModelForSequenceClassification.from_pretrained(model_dir, dtype=torch.float16, device_map="auto")
     model.eval()
     texts = [
         "This is a test sentence.",