Update app.py

app.py

@@ -2,11 +2,9 @@ import gradio as gr
 import torch
 from peft import PeftModel, PeftConfig, LoraConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from datasets import load_dataset
-from trl import SFTTrainer
 # import torch
 from transformers import StoppingCriteria, AutoModelForCausalLM, AutoTokenizer, StoppingCriteriaList
-
+import numpy as np
 
 ref_model = AutoModelForCausalLM.from_pretrained("w601sxs/b1ade-1b", torch_dtype=torch.bfloat16)
 
@@ -29,14 +27,9 @@ stop_ids = [tokenizer.encode(w)[0] for w in stop_words]
 stop_criteria = KeywordsStoppingCriteria(stop_ids)
 
 
-
-
-
-import numpy as np
-
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
-    model.config.pad_token_id = model.config.eos_token_id
+    ref_model.config.pad_token_id = ref_model.config.eos_token_id
 
 # Define your color-coding labels; if prob > x, then label = y; Sorted in descending probability order!
 probs_to_label = [
@@ -46,7 +39,6 @@ probs_to_label = [
     (0.5, "50%"),
     (0.1, "10%"),
     (0.01, "1%"),
-    
 ]
 
 
@@ -78,7 +70,7 @@ def get_tokens_and_labels(prompt):
 
     # Important: you might need to find a tokenization character to replace (e.g. "Ġ" for BPE) and get the correct
     # spacing into the final output 👼
-    if model.config.is_encoder_decoder:
+    if ref_model.config.is_encoder_decoder:
         highlighted_out = []
     else:
         input_tokens = tokenizer.convert_ids_to_tokens(inputs.input_ids[0])