Update app.py

app.py

@@ -1,15 +1,8 @@
 import gradio as gr
+import spaces
 from peft import AutoPeftModelForCausalLM
 from transformers import AutoTokenizer
 
-# Load the model and tokenizer
-load_in_4bit = True  # Adjust based on your setup
-model = AutoPeftModelForCausalLM.from_pretrained(
-    "shakaryan/lebedev_qwen2.5",
-    load_in_4bit=load_in_4bit,
-)
-tokenizer = AutoTokenizer.from_pretrained("shakaryan/lebedev_qwen2.5")
-
 # Define the prompt template
 alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
 
@@ -22,10 +15,17 @@ alpaca_prompt = """Below is an instruction that describes a task, paired with an
 ### Response:
 {}"""
 
-EOS_TOKEN = tokenizer.eos_token  # Ensure proper sequence termination
-
 # Function to generate responses
+@spaces.GPU
 def generate_response(input_text, instruction_text):
+    # Load the model and tokenizer within the GPU context
+    model = AutoPeftModelForCausalLM.from_pretrained(
+        "shakaryan/lebedev_qwen2.5",
+        load_in_4bit=True,  # Adjust based on your setup
+    ).to("cuda")
+    tokenizer = AutoTokenizer.from_pretrained("shakaryan/lebedev_qwen2.5")
+    EOS_TOKEN = tokenizer.eos_token  # Ensure proper sequence termination
+
     # Format the prompt
     formatted_prompt = alpaca_prompt.format(instruction_text, input_text, "") + EOS_TOKEN
     # Tokenize and generate response
@@ -46,7 +46,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         instruction_text = gr.Textbox(
             label="Instruction Text",
-            value="You are a blogger named Artemiy Lebedev, your purpose is to generate a post on Russian based on the post article",
+            value="You are a blogger named Artemiy Lebedev, your purpose is to generate a post in Russian based on the post article",
             lines=3,
         )
     with gr.Row():