move to cpu

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 import re
-import torch
 import contractions
 import pandas as pd
 from transformers import BartTokenizer, BartForConditionalGeneration
@@ -77,9 +76,8 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 # Load model and tokenizer
-device = torch.device("cpu") 
 MODEL_PATH = "./models/fine-tuned_bart_base"
-model = BartForConditionalGeneration.from_pretrained(MODEL_PATH).to(device)
+model = BartForConditionalGeneration.from_pretrained(MODEL_PATH).cpu()
 tokenizer = BartTokenizer.from_pretrained(MODEL_PATH)
 
 #Helper functions
@@ -279,6 +277,7 @@ def summarize_text(txt):
   txt = preprocess_text(txt)
   txt = anonymize_speakers(txt, speaker_1, speaker_2)
   inputs = tokenizer(txt, return_tensors="pt")
+  inputs = {k: v.cpu() for k, v in inputs.items()}
   summary_ids = model.generate(**inputs)
   summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
   summary = deanonymize_speakers(summary, speaker_1, speaker_2)