Update app.py

app.py

@@ -1,15 +1,3 @@
-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
@@ -29,44 +17,21 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens,
         return_tensors="pt",  # Return tensors for PyTorch
     )
 
-    # Debug inputs
-    print("Inputs:", inputs)
-
-    # Extract input_ids correctly
-    input_ids = inputs["input_ids"]  # Ensure the correct field is used
+    # Ensure input_ids is extracted
+    input_ids = inputs  # Tensor provided directly
     print("Input IDs shape:", input_ids.shape)
 
     # Generate response
     output = model.generate(
-        input_ids=input_ids,
+        input_ids=input_ids,  # Pass tensor explicitly as input_ids
         max_new_tokens=max_tokens,
         use_cache=True,
         temperature=temperature,
         min_p=min_p,
     )
 
-    # Debug output
-    print("Generated Output Shape:", output.shape)
-    print("Generated Output:", output)
-
     # Decode and format the response
-    # Use `decode` for the first sequence in the batch
     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()