Update app.py

app.py

@@ -1,47 +1,38 @@
 import gradio as gr
-import os
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
 
-# 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)
+# Initialize model
 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)
 
-# Define prompt structure
-@spaces.GPU  # This will handle spaces for either GPU or CPU as available
+@spaces.GPU
 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)
+    # 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)
 
     # Generate model response
     with torch.no_grad():
-        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()
+        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()
 
-    return response  # Return only bot's response for display in the right column
+    return response
 
-# Set up Gradio interface
+# Gradio interface setup
 with gr.Blocks() as demo:
     gr.Markdown("# Clinical Trial Chatbot")
 
     with gr.Row():
-        # Left column for inputs
+        # Left sidebar for inputs
         with gr.Column():
             system_instruction = gr.Textbox(
                 placeholder="Enter system instruction here...", label="System Instruction")
@@ -49,12 +40,12 @@ with gr.Blocks() as demo:
                 placeholder="Type your message here...", label="Your Message")
             submit_btn = gr.Button("Submit")
 
-        # Right column for displaying bot's response
+        # Right column for displaying bot response
         with gr.Column():
             response_display = gr.Textbox(
                 label="Bot Response", interactive=False, placeholder="Response will appear here.")
 
-    # Link the submit button to the generate_response function
+    # Link submit button to the generate_response function
     submit_btn.click(generate_response, [system_instruction, user_input], response_display)
 
 # Launch the app