amy011872
/

LawToken-7B-a2

Text Generation

Generated from Trainer

text-generation-inference

Inference Endpoints

Model card Files Files and versions Community

amy011872 commited on 21 days ago

Commit

41d0d3e

•

1 Parent(s): 7a1ba7d

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -106,7 +106,7 @@ class EndpointHandler():
             outputs = self.model(**inputs)
         outputs_logits = outputs.logits[0, -1, self.law_token_ids]
-        base_input = tokenizer("<cite>", return_tensors="pt").to("cuda")
         with torch.no_grad():
             base_output = self.model(**base_input)

             outputs = self.model(**inputs)
         outputs_logits = outputs.logits[0, -1, self.law_token_ids]
+        base_input = self.tokenizer("<cite>", return_tensors="pt").to("cuda")
         with torch.no_grad():
             base_output = self.model(**base_input)