Update app.py

app.py

@@ -5,6 +5,10 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 # import torch
 from transformers import StoppingCriteria, AutoModelForCausalLM, AutoTokenizer, StoppingCriteriaList
 import numpy as np
+import spacy
+from spacy import displacy
+from spacy.tokens import Span
+from spacy.tokens import Doc
 
 ref_model = AutoModelForCausalLM.from_pretrained("w601sxs/b1ade-1b", torch_dtype=torch.bfloat16)
 
@@ -48,7 +52,7 @@ def get_tokens_and_labels(prompt):
     """
     Given the prompt (text), return a list of tuples (decoded_token, label)
     """
-    inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
+    inputs = tokenizer([prompt], return_tensors="pt")
     outputs = ref_model.generate(
         **inputs,
         max_new_tokens=1000,
@@ -57,7 +61,7 @@ def get_tokens_and_labels(prompt):
         stopping_criteria=StoppingCriteriaList([stop_criteria])
     )
     # Important: don't forget to set `normalize_logits=True` to obtain normalized probabilities (i.e. sum(p) = 1)
-    transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)
+    transition_scores = ref_model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)
     transition_proba = np.exp(transition_scores.double().cpu())
     
     # print(transition_proba)
@@ -89,47 +93,6 @@ def get_tokens_and_labels(prompt):
 
 
 
-import spacy
-from spacy import displacy
-from spacy.tokens import Span
-from spacy.tokens import Doc
-
-
-
-def render_output(prompt):
-    output = get_tokens_and_labels(prompt)
-    nlp = spacy.blank("en")
-    doc = nlp(''.join([a[0] for a in output]).replace('Ġ',' ').replace('Ċ','\n'))
-    words = [a[0].replace('Ġ',' ').replace('Ċ','\n') for a in output]#[:indices[2]]
-    doc = Doc(nlp.vocab, words=words)
-
-    doc.spans["sc"]=[]
-    c = 0
-
-    for outs in output:
-        tmpouts = outs[0].replace('Ġ','').replace('Ċ','\n')
-        # print(c, "to", c+len(tmpouts)," : ", tmpouts)
-
-        if outs[1] is not None:
-            doc.spans["sc"].append(Span(doc, c, c+1, outs[1] ))  
-
-        c+=1
-
-    #     if c>indices[2]-1:
-    #         break
-
-
-    options = {'colors' : {
-            '99%': '#44ce1b',
-            '95%': '#bbdb44',
-            '90%': '#f7e379',
-            '50%': '#fec12a',
-            '10%': '#f2a134',
-            '1%': '#e51f1f',
-            '': '#e51f1f',
-    }}
-
-    return displacy.render(doc, style="span", options = options)
 
 
 
@@ -144,7 +107,7 @@ def predict(text):
     
 
 demo = gr.Interface(
-    fn=render_output, 
+    fn=get_tokens_and_labels, 
     inputs='text',
     outputs='text',
 )