Update app.py

app.py

@@ -4,17 +4,28 @@ from arabert.aragpt2.grover.modeling_gpt2 import GPT2LMHeadModel
 from transformers import AutoTokenizer
 import re
 import torch
+import os
 
-# Check if CUDA is available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+# 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")
 
 model_name = "Naseej/AskMe-Large"
 tokenizer = AutoTokenizer.from_pretrained(model_name, bos_token='<|startoftext|>',
                                          eos_token='<|endoftext|>', pad_token='<|pad|>')
-model = GPT2LMHeadModel.from_pretrained(model_name).to(device)  # Move model to GPU
+model = GPT2LMHeadModel.from_pretrained(model_name)
 model.resize_token_embeddings(len(tokenizer))
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+
+# Configure pipeline with the device parameter for ZeroGPU
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
 
 def generate_response(message, history, num_beams=4, temperature=0.99, do_sample=True, top_k=60, top_p=0.9):
     prompt = f'Prompt: {message}\nAnswer:'
@@ -94,4 +105,8 @@ with gr.Blocks(css=css) as demo:
     
     clear_btn.click(lambda: None, None, chatbot, queue=False)
 
-demo.launch()
+# 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()