Update app.py

app.py

@@ -4,6 +4,7 @@ from clean_data import cleaned_complaints
 import numpy as np
 from scipy.special import softmax
 import gradio as gr
+from transformers import PegasusForConditionalGeneration
 
 # Preprocess text (username and link placeholders)
 def preprocess(text):
@@ -16,7 +17,8 @@ def preprocess(text):
 
 # load model
 MODEL = f"ThirdEyeData/Complaints_Roberta"
-model = AutoModelForSequenceClassification.from_pretrained(MODEL)
+model = PegasusForConditionalGeneration.from_pretrained(MODEL)
+#model = AutoModelForSequenceClassification.from_pretrained(MODEL)
 #model.save_pretrained(MODEL)