Update app.py

app.py

@@ -4,75 +4,58 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
 
-# Initialize model
+# Check if we're running in a Hugging Face Space with zero GPU constraints
+IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
+IS_SPACE = os.environ.get("SPACE_ID", None) is not None
+
+# Determine device (set to CPU for zero-GPU)
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# Load model and tokenizer
 model_name = "linjc16/Panacea-7B-Chat"
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
 model.to(device)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Store prompts
-saved_prompts = {}
-
-@spaces.GPU
-def generate_response(system_instruction, user_input, prompt_name):
-    # Format the prompt using the messages structure
-    messages = [
-        {"role": "system", "content": system_instruction},
-        {"role": "user", "content": user_input},
-    ]
-    encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt").to(device)
-    model_inputs = encodeds.to(device)
+# Define prompt structure
+@spaces.GPU  # This will handle spaces for either GPU or CPU as available
+def generate_response(system_instruction, user_input):
+    # Format the prompt with the system instruction and user input
+    prompt = f"{system_instruction}\n\nUser: {user_input}\nBot:"
+    
+    # Tokenize and prepare inputs
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
     # Generate model response
     with torch.no_grad():
-        generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0].split("Bot:")[-1].strip()
-
-    # Save the prompt with the specified name
-    saved_prompts[prompt_name] = {"system": system_instruction, "user": user_input}
-    return response, list(saved_prompts.keys())
-
-def load_prompt(prompt_name):
-    # Load selected prompt into the input fields
-    if prompt_name in saved_prompts:
-        prompt = saved_prompts[prompt_name]
-        return prompt["system"], prompt["user"]
-    return "", ""
+        outputs = model.generate(**inputs, max_new_tokens=100, do_sample=True)
+    
+    # Decode response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Bot:")[-1].strip()
 
-def rename_prompt(old_name, new_name):
-    # Rename a saved prompt
-    if old_name in saved_prompts:
-        saved_prompts[new_name] = saved_prompts.pop(old_name)
-    return list(saved_prompts.keys())
+    return response  # Return only bot's response for display in the right column
 
-# Gradio interface setup
+# Set up Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Clinical Trial Chatbot with Prompt Saving")
+    gr.Markdown("# Clinical Trial Chatbot")
 
     with gr.Row():
-        # Left sidebar for saved prompts and inputs
+        # Left column for inputs
         with gr.Column():
-            saved_prompts_list = gr.Dropdown(label="Saved Prompts", choices=[])
-            rename_input = gr.Textbox(label="Rename Prompt")
-            rename_button = gr.Button("Rename")
             system_instruction = gr.Textbox(
                 placeholder="Enter system instruction here...", label="System Instruction")
             user_input = gr.Textbox(
                 placeholder="Type your message here...", label="Your Message")
-            prompt_name = gr.Textbox(placeholder="Enter prompt name...", label="Prompt Name")
-            submit_btn = gr.Button("Save & Submit")
+            submit_btn = gr.Button("Submit")
 
-        # Right column for displaying bot response
+        # Right column for displaying bot's response
         with gr.Column():
             response_display = gr.Textbox(
                 label="Bot Response", interactive=False, placeholder="Response will appear here.")
 
-    # Link actions to functions
-    submit_btn.click(generate_response, [system_instruction, user_input, prompt_name], 
-                     [response_display, saved_prompts_list])
-    saved_prompts_list.change(load_prompt, saved_prompts_list, [system_instruction, user_input])
-    rename_button.click(rename_prompt, [saved_prompts_list, rename_input], saved_prompts_list)
+    # Link the submit button to the generate_response function
+    submit_btn.click(generate_response, [system_instruction, user_input], response_display)
 
 # Launch the app
-demo.launch()
+demo.launch(share=True)