Update app.py

app.py

@@ -1,11 +1,10 @@
 import streamlit as st
-from flask import Flask, request, render_template, redirect
-from transformers import BartTokenizer, BartForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 # preprocess input
 # return input_ids matrix
-tokenizer = BartTokenizer.from_pretrained("sshleifer/distilbart-cnn-6-6")
-model = BartForConditionalGeneration.from_pretrained("sshleifer/distilbart-cnn-6-6")
+tokenizer = AutoTokenizer.from_pretrained("SophieTr/fine-tune-Pegasus")
+model = AutoModelForSeq2SeqLM.from_pretrained("SophieTr/fine-tune-Pegasus")
 
 def preprocess(inp):
     input_ids = tokenizer(inp, return_tensors="pt").input_ids