poiccard
/

jais-13b-chat-adn

Text Generation

Inference Endpoints

Model card Files Files and versions Community

poiccard commited on Sep 21, 2023

Commit

91174fb

•

1 Parent(s): 4f6f950

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -12,7 +12,7 @@ class EndpointHandler():
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", trust_remote_code=True)
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:

         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", offload_folder='offload', trust_remote_code=True)
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: