Update app.py

app.py

@@ -4,19 +4,7 @@ from arabert.aragpt2.grover.modeling_gpt2 import GPT2LMHeadModel
 from transformers import AutoTokenizer
 import re
 import torch
-import os
-
-# For ZeroGPU on Hugging Face
-# Checking if we're on Hugging Face infrastructure
-HF_SPACE = os.environ.get("SPACE_ID") is not None
-
-# Get appropriate device for ZeroGPU
-if torch.cuda.is_available() and HF_SPACE:
-    device = 0  # For ZeroGPU, use device index 0
-    print("Using ZeroGPU on Hugging Face")
-else:
-    device = -1  # CPU fallback
-    print("Using CPU")
+import spaces  # Import the spaces module for ZeroGPU
 
 model_name = "Naseej/AskMe-Large"
 tokenizer = AutoTokenizer.from_pretrained(model_name, bos_token='<|startoftext|>',
@@ -24,10 +12,16 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, bos_token='<|startoftext|>
 model = GPT2LMHeadModel.from_pretrained(model_name)
 model.resize_token_embeddings(len(tokenizer))
 
-# Configure pipeline with the device parameter for ZeroGPU
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
+# For ZeroGPU, we'll move the model to CUDA inside the decorated function
+# Create the generator pipeline without specifying device
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
+# ZeroGPU-decorated function for text generation
+@spaces.GPU(duration=60)  # Set duration based on your needs
 def generate_response(message, history, num_beams=4, temperature=0.99, do_sample=True, top_k=60, top_p=0.9):
+    # Move model to CUDA inside the decorated function
+    generator.model = generator.model.to('cuda')
+    
     prompt = f'Prompt: {message}\nAnswer:'
     pred_text = generator(prompt,
                          pad_token_id=tokenizer.eos_token_id,
@@ -45,6 +39,8 @@ def generate_response(message, history, num_beams=4, temperature=0.99, do_sample
     except:
         pred_sentiment = "لم أستطع توليد إجابة. يرجى إعادة صياغة السؤال."
 
+    # Move model back to CPU to free GPU memory
+    generator.model = generator.model.to('cpu')
     return pred_sentiment
 
 # Properly format the chat message handler
@@ -105,8 +101,4 @@ with gr.Blocks(css=css) as demo:
     
     clear_btn.click(lambda: None, None, chatbot, queue=False)
 
-# For ZeroGPU on Hugging Face, we need to use specific launch parameters
-if HF_SPACE:
-    demo.launch(share=False, server_name="0.0.0.0")
-else:
-    demo.launch()
+demo.launch()