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").to("cuda") # Move model to CUDA 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 messages = [{"role": "system", "content": system_message}] for val in history: if val[0]: messages.append({"role": "user", "content": val[0]}) if val[1]: messages.append({"role": "assistant", "content": val[1]}) messages.append({"role": "user", "content": message}) # Tokenize the input messages inputs = tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", # Return tensors for PyTorch ) # Ensure input_ids is moved to the same device as the model input_ids = inputs.to("cuda") # Move input_ids to CUDA print("Input IDs shape:", input_ids.shape) # Generate response output = model.generate( 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 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()