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(