Update app.py

app.py

@@ -3,13 +3,18 @@ from transformers import pipeline
 from arabert.aragpt2.grover.modeling_gpt2 import GPT2LMHeadModel
 from transformers import AutoTokenizer
 import re
+import torch
+
+# Check if CUDA is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 
 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)
+model = GPT2LMHeadModel.from_pretrained(model_name).to(device)  # Move model to GPU
 model.resize_token_embeddings(len(tokenizer))
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
 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:'