Add cpu inference option for testing

app.py

@@ -4,10 +4,15 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, Stopping
 import time
 import numpy as np
 from torch.nn import functional as F
-
-
-m = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-7b", torch_dtype=torch.float16).cuda()
-tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b")
+import os
+token_key = os.environ.get(“HUGGING_FACE_HUB_TOKEN”)
+
+if torch.cuda.is_available():
+    m = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=token_key, torch_dtype=torch.float16).cuda()
+    tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=token_key)
+else:
+    m = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=token_key, torch_dtype=torch.float16)
+    tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=token_key)
 generator = pipeline('text-generation', model=m, tokenizer=tok, device=0)
 
 
@@ -29,8 +34,12 @@ class StopOnTokens(StoppingCriteria):
 
 def contrastive_generate(text, bad_text):
     with torch.no_grad():
-        tokens = tok(text, return_tensors="pt")['input_ids'].cuda()[:,:4096-1024]
-        bad_tokens = tok(bad_text, return_tensors="pt")['input_ids'].cuda()[:,:4096-1024]
+        if torch.cuda_is_available():
+            tokens = tok(text, return_tensors="pt")['input_ids'].cuda()[:,:4096-1024]
+            bad_tokens = tok(bad_text, return_tensors="pt")['input_ids'].cuda()[:,:4096-1024]
+        else:
+            tokens = tok(text, return_tensors="pt")['input_ids'][:,:4096-1024]
+            bad_tokens = tok(bad_text, return_tensors="pt")['input_ids'][:,:4096-1024]
         history = None
         bad_history = None
         curr_output = list()
@@ -83,12 +92,13 @@ def system_update(msg):
 
 
 with gr.Blocks() as demo:
+    gr.Markdown("###StableLM-tuned-Alpha-7B Chat")
     with gr.Row():
         with gr.Column():
             chatbot = gr.Chatbot([])
             clear = gr.Button("Clear")
         with gr.Column():
-            system_msg = gr.Textbox(start_message, label="System Message", interactive=True)
+            system_msg = start_message#gr.Textbox(start_message, label="System Message", interactive=True)
             msg = gr.Textbox(label="Chat Message")
 
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(