Update app.py

app.py

@@ -54,19 +54,17 @@ def predict(message, history, system_prompt, temperature, max_tokens):
     outputs = []
     for text in streamer:
         outputs.append(text)
-        stop = False
-        for tok in stop_tokens:
-            if tok in text:
-                stop = True
-                break
-        if stop:
-            yield "".join(outputs)
+        if text in stop_tokens:
             break
+        print(text)
+        yield "".join(outputs)
+
+
 
 if __name__ == "__main__":
     args = parse_args()
-    tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-12b-chat", trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-2-12b-chat", trust_remote_code=True, torch_dtype=torch.bfloat16)
+    tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-chat", trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-2-chat", trust_remote_code=True, torch_dtype=torch.bfloat16)
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model = model.to(device)
     gr.ChatInterface(
@@ -85,4 +83,4 @@ if __name__ == "__main__":
             gr.Slider(100, 2048, 1024, label="Max Tokens"),
         ],
         additional_inputs_accordion_name="Parameters",
-    ).queue().launch()
+    ).queue().launch()