update app.py

app.py

@@ -21,8 +21,10 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 from threading import Thread
 
-tokenizer = AutoTokenizer.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1")
-model = AutoModelForCausalLM.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1", torch_dtype=torch.float16)
+#tokenizer = AutoTokenizer.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1")
+#model = AutoModelForCausalLM.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1", torch_dtype=torch.float16)
+model = AutoModelForCausalLM.from_pretrained("mattshumer/mistral-8x7b-chat", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("mattshumer/mistral-8x7b-chat")
 model = model.to('cuda:0')
 
 class StopOnTokens(StoppingCriteria):
@@ -41,6 +43,10 @@ def predict(message, history):
     messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])  #curr_system_message +
                 for item in history_transformer_format])
 
+    #     x = tok.encode(PROMPT, return_tensors="pt").cuda()
+    #     x = model.generate(x, max_new_tokens=512).cpu()
+    #     return tok.batch_decode(x)
+
     model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(