median-group
/

Meta-Llama-3-8B-logits

Text Generation

text-generation-inference

Inference Endpoints

Model card Files Files and versions Community

HazardSpence commited on Oct 23, 2024

Commit

ef3cfb9

·

verified ·

1 Parent(s): 9151243

Update logits.py

Files changed (1) hide show

logits.py +1 -1

logits.py CHANGED Viewed

@@ -11,7 +11,7 @@ class LogitsPredictor:
     def setup(self, model_path="./"):
         """Load the model into memory to make running multiple predictions efficient"""
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-        self.model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
         self.separator = list(filter(lambda x: x != self.tokenizer.bos_token_id, self.tokenizer.encode("\n")))[0]
     def quantiles(self, sorted_probs, ranks):

     def setup(self, model_path="./"):
         """Load the model into memory to make running multiple predictions efficient"""
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path)
         self.separator = list(filter(lambda x: x != self.tokenizer.bos_token_id, self.tokenizer.encode("\n")))[0]
     def quantiles(self, sorted_probs, ranks):