Update app.py

app.py

@@ -2,19 +2,18 @@ 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
-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)
+ref_model = ref_model.to('cuda')
+ref_model.eval()
 
 tokenizer = AutoTokenizer.from_pretrained("w601sxs/b1ade-1b")
 
-ref_model.eval()
 
 class KeywordsStoppingCriteria(StoppingCriteria):
     def __init__(self, keywords_ids:list):
@@ -30,6 +29,7 @@ stop_words = ['>', ' >','> ']
 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
@@ -43,16 +43,14 @@ probs_to_label = [
     (0.5, "50%"),
     (0.1, "10%"),
     (0.01, "1%"),
+    
 ]
-
-
-
-
+import numpy as np
 def get_tokens_and_labels(prompt):
     """
     Given the prompt (text), return a list of tuples (decoded_token, label)
     """
-    inputs = tokenizer([prompt], return_tensors="pt")
+    inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
     outputs = ref_model.generate(
         **inputs,
         max_new_tokens=1000,
@@ -91,10 +89,45 @@ def get_tokens_and_labels(prompt):
 
     return highlighted_out
 
+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)
 
 
 def predict(text):
@@ -107,9 +140,9 @@ def predict(text):
     
 
 demo = gr.Interface(
-    fn=get_tokens_and_labels, 
+    fn=render_output, 
     inputs='text',
-    outputs='text',
+    outputs='html',
 )
 
 demo.launch()