Update app.py

app.py

@@ -5,7 +5,8 @@ import time
 import numpy as np
 from torch.nn import functional as F
 import os
-# auth_key = os.environ["HF_ACCESS_TOKEN"]
+from threading import Thread
+
 print(f"Starting to load the model to memory")
 m = AutoModelForCausalLM.from_pretrained(
     "stabilityai/stablelm-tuned-alpha-7b", torch_dtype=torch.float16).cuda()
@@ -28,18 +29,6 @@ class StopOnTokens(StoppingCriteria):
                 return True
         return False
 
-def streaming_generate(text, bad_text=None):
-    
-    return model_output
-    
-    
-def generate(text, bad_text=None):
-    stop = StopOnTokens()
-    result = generator(text, max_new_tokens=1024, num_return_sequences=1, num_beams=1, do_sample=True,
-                       temperature=1.0, top_p=0.95, top_k=1000, stopping_criteria=StoppingCriteriaList([stop]))
-    return result[0]["generated_text"].replace(text, "")
-
-
 def user(user_message, history):
     history = history + [[user_message, ""]]
     return "", history, history
@@ -51,8 +40,8 @@ def bot(history, curr_system_message):
         "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]])
                 for item in history])
     
-    model_inputs = tok(messages, return_tensors="pt")['input_ids'].cuda()[:, :4096-1024]
-
+    #model_inputs = tok([messages], return_tensors="pt")['input_ids'].cuda()[:, :4096-1024]
+    model_inputs = tok([messages], return_tensors="pt").to("cuda")
     streamer = TextIteratorStreamer(tok, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         model_inputs,
@@ -68,10 +57,11 @@ def bot(history, curr_system_message):
     t = Thread(target=m.generate, kwargs=generate_kwargs)
     t.start()
 
-    model_output = ""
+    print(history)
     for new_text in streamer:
+        print(new_text)
         history[-1][1] += new_text
-        yield history
+        yield history, history
         
     return history, history