Upload code/ with huggingface_hub

code/inference.py

@@ -1,7 +1,7 @@
 import logging, requests, os, io, glob, time
 import json
 
-
+from transformers import  T5TokenizerFast
 from transformers import BertTokenizer
 from transformers import PreTrainedModel
 import torch
@@ -276,10 +276,12 @@ JSON_CONTENT_TYPE = 'application/json'
     
 
 # loads the model into memory from disk and returns it
-def model_fn():
-    
-    model = AutoModelForSeq2SeqLM.from_pretrained("canovich/myprivateee")
-    return model
+def model_fn(model_dir):
+  # Load model from HuggingFace Hub
+  tokenizer = T5TokenizerFast(model_dir, extra_ids=0)
+  model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+  return model, tokenizer
+
 
 
 # Perform prediction on the deserialized object, with the loaded model