Update app.py

app.py

@@ -1,3 +1,15 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Prevent CUDA initialization outside ZeroGPU
+
+import spaces  # Import spaces first
+import gradio as gr
+from peft import AutoPeftModelForCausalLM
+from transformers import AutoTokenizer
+
+# Load the model and tokenizer globally
+model = AutoPeftModelForCausalLM.from_pretrained("eforse01/lora_model")
+tokenizer = AutoTokenizer.from_pretrained("eforse01/lora_model")
+
 @spaces.GPU(duration=120)  # Decorate the function for ZeroGPU
 def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, min_p):
     # Construct messages for the chat template
@@ -18,7 +30,7 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens,
     )
 
     # Ensure input_ids is extracted
-    input_ids = inputs  # Tensor provided directly
+    input_ids = inputs  # Directly using tensor returned from apply_chat_template
     print("Input IDs shape:", input_ids.shape)
 
     # Generate response
@@ -30,8 +42,27 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens,
         min_p=min_p,
     )
 
+    # Debug output
+    print("Generated Output Shape:", output.shape)
+    print("Generated Output:", output)
+
     # Decode and format the response
     response = tokenizer.decode(output[0], skip_special_tokens=True)
 
     # Yield the response
     yield response.split("assistant")[-1]
+
+
+# Gradio Interface
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=2048, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=1.5, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.99, step=0.01, label="Min-p"),
+    ],
+)
+
+if __name__ == "__main__":
+    demo.launch()