Update app.py

app.py

@@ -1,25 +1,27 @@
-import gradio as gr
-from transformers import BartTokenizer, BartForConditionalGeneration
-
-# Load your model and tokenizer from the Hugging Face Hub
-model_name = "bart_headline_model"  # Replace with your uploaded model name
-tokenizer = BartTokenizer.from_pretrained(model_name)
-model = BartForConditionalGeneration.from_pretrained(model_name)
-
-# Define the function for generating captions
-def generate_caption(article):
-    inputs = tokenizer(article, return_tensors="pt", max_length=128, truncation=True)
-    outputs = model.generate(inputs["input_ids"], max_length=20, num_beams=5, early_stopping=True)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Create the Gradio interface
-interface = gr.Interface(
-    fn=generate_caption,
-    inputs="text",
-    outputs="text",
-    title="Headline Generator",
-    description="Enter an article to generate a headline.",
-)
-
-# Launch the app
-interface.launch()
+import gradio as gr
+from transformers import BartTokenizer, BartForConditionalGeneration
+
+# Load your model and tokenizer from the Hugging Face Hub
+model_name = "bart_headline_model"  # Replace with your uploaded model name
+tokenizer = BartTokenizer.from_pretrained(model_name)
+#model = BartForConditionalGeneration.from_pretrained(model_name)
+model = BartForConditionalGeneration.from_pretrained(model_name, device_map="cpu")
+
+
+# Define the function for generating captions
+def generate_caption(article):
+    inputs = tokenizer(article, return_tensors="pt", max_length=128, truncation=True)
+    outputs = model.generate(inputs["input_ids"], max_length=20, num_beams=5, early_stopping=True)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Create the Gradio interface
+interface = gr.Interface(
+    fn=generate_caption,
+    inputs="text",
+    outputs="text",
+    title="Headline Generator",
+    description="Enter an article to generate a headline.",
+)
+
+# Launch the app
+interface.launch()