Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from datasets import load_dataset
-from transformers import BartForQuestionAnswering, BartTokenizer
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
 import torch
 import logging
 from fastapi import FastAPI, HTTPException
@@ -19,9 +19,11 @@ logger.info("Dataset loaded successfully")
 # Load model and tokenizer
 logger.info("Loading the model and tokenizer")
 model_name = "facebook/bart-large-cnn"
-tokenizer = BartTokenizer.from_pretrained(model_name)
-model = BartForQuestionAnswering.from_pretrained(model_name)
-logger.info("Model and tokenizer loaded successfully")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+logger.info(f"Model and tokenizer loaded successfully. Using device: {device}")
 
 def clean_answer(answer):
     special_tokens = set(tokenizer.all_special_tokens)
@@ -51,7 +53,7 @@ def answer_question(question):
                 padding='max_length'
             )
             
-            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+            inputs = {k: v.to(device) for k, v in inputs.items()}
             
             logger.info(f"Input tokens shape: {inputs['input_ids'].shape}")
             
@@ -131,7 +133,7 @@ iface = gr.Interface(
 # Mount Gradio app to FastAPI
 app = gr.mount_gradio_app(app, iface, path="/")
 
-# For local development and testing
+# For Hugging Face Spaces
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)