Update app.py

app.py

@@ -7,18 +7,16 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 model_repo_id = "Ajay12345678980/QA_bot"  # Replace with your model repository ID
 
 # Initialize the model and tokenizer
-model = GPT2LMHeadModel.from_pretrained(model_repo_id)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = GPT2LMHeadModel.from_pretrained(model_repo_id).to(device)
 tokenizer = GPT2Tokenizer.from_pretrained(model_repo_id)
 
 # Define the prediction function
 def generate_answer(question):
-    input_ids = tokenizer.encode(question, return_tensors="pt").to("cuda")
-
-    # Create the attention mask and pad token id
-    attention_mask = torch.ones_like(input_ids).to("cuda")
+    input_ids = tokenizer.encode(question, return_tensors="pt").to(device)
+    attention_mask = torch.ones_like(input_ids).to(device)
     pad_token_id = tokenizer.eos_token_id
 
-    #output = model[0].generate(
     output = model.generate(
         input_ids,
         max_new_tokens=100,
@@ -30,17 +28,14 @@ def generate_answer(question):
     start_index = decoded_output.find("Answer")
     end_index = decoded_output.find("<ANSWER_ENDED>")
 
-    if end_index != -1:
-        # Extract the text between "Answer" and "<ANSWER_ENDED>"
-        answer_text = decoded_output[start_index + len("Answer"):end_index].strip()
+    if start_index != -1:
+        if end_index != -1:
+            answer_text = decoded_output[start_index + len("Answer"):end_index].strip()
+        else:
+            answer_text = decoded_output[start_index + len("Answer"):].strip()
         return answer_text
     else:
-        # If "<ANSWER_ENDED>" is not found, return the text following "Answer"
-        answer_text = decoded_output[start_index + len("Answer"):].strip()
-        return answer_text
-
-    #return tokenizer.decode(output[0], skip_special_tokens=True)
-    #return tokenizer.decode(output, skip_special_tokens=True)
+        return "Sorry, I couldn't generate an answer."
 
 # Gradio interface setup
 interface = gr.Interface(
@@ -48,7 +43,7 @@ interface = gr.Interface(
     inputs="text", 
     outputs="text", 
     title="GPT-2 Text Generation",
-    description="Enter some text and see what the model generates!"
+    description="Enter a question and see what the model generates!"
 )
 
 # Launch the Gradio app