Update app.py

app.py

@@ -10,9 +10,27 @@ model_path = 'THUDM/BPO'
 
 device = 'cuda'
 
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, add_prefix_space=True)
-model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, load_in_8bit=True)
-model = model.eval()
+if torch.cuda.is_available():
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, add_prefix_space=True)
+    model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, device_map=device, load_in_8bit=True)
+    model = model.eval()
+
+
+DESCRIPTION = """This Space demonstrates model [BPO](https://huggingface.co/THUDM/BPO), which is built on LLaMA-2-7b-chat.
+BPO aims to improve the alignment of LLMs with human preferences by optimizing user prompts.
+
+Feel free to play with it, or duplicate to run generations without a queue! 🔎 For more details about the BPO model, take a look [at our paper](https://arxiv.org/pdf/2311.04155.pdf).
+"""
+
+LICENSE = """
+---
+As BPO is a fine-tuned version of [Llama-2-7b-chat](https://huggingface.co/meta-llama/Llama-2-7b-chat) by Meta,
+this demo is governed by the original [license](https://huggingface.co/spaces/CCCCCC/BPO_demo/blob/main/LICENSE.txt) and [acceptable use policy](https://huggingface.co/spaces/CCCCCC/BPO_demo/blob/main/USE_POLICY.md).
+"""
+
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
 prompt_template = "[INST] You are an expert prompt engineer. Please help me improve this prompt to get a more helpful and harmless response:\n{} [/INST]"
 
@@ -119,9 +137,13 @@ def reset_state():
 def update_textbox_from_dropdown(selected_example):
     return selected_example
 
-with gr.Blocks() as demo:
+with gr.Blocks("sty.css") as demo:
     gr.HTML("""<h1 align="center">Prompt Preference Optimizer</h1>""")
 
+    gr.Markdown(DESCRIPTION)
+
+    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
+    
     chatbot = gr.Chatbot(label="Prompt Optimization Chatbot")
     with gr.Row():
         with gr.Column(scale=4):
@@ -137,6 +159,8 @@ with gr.Blocks() as demo:
             top_p = gr.Slider(0, 1, value=0.9, step=0.01, label="Top P", interactive=True)
             temperature = gr.Slider(0, 1, value=0.6, step=0.01, label="Temperature", interactive=True)
 
+    gr.Markdown(LICENSE)
+    
     dropdown.change(update_textbox_from_dropdown, dropdown, user_input)
     
     history = gr.State([])