Update app.py

app.py

@@ -7,6 +7,13 @@ from trl import SFTTrainer
 # import torch
 from transformers import StoppingCriteria, AutoModelForCausalLM, AutoTokenizer, StoppingCriteriaList
 
+
+ref_model = AutoModelForCausalLM.from_pretrained("w601sxs/b1ade-1b", torch_dtype=torch.bfloat16)
+
+tokenizer = AutoTokenizer.from_pretrained("w601sxs/b1ade-1b")
+
+ref_model.eval()
+
 class KeywordsStoppingCriteria(StoppingCriteria):
     def __init__(self, keywords_ids:list):
         self.keywords = keywords_ids
@@ -42,11 +49,7 @@ probs_to_label = [
     
 ]
 
-ref_model = AutoModelForCausalLM.from_pretrained("w601sxs/b1ade-1b", torch_dtype=torch.bfloat16)
 
-tokenizer = AutoTokenizer.from_pretrained("w601sxs/b1ade-1b")
-
-ref_model.eval()
 
 
 def get_tokens_and_labels(prompt):