Update app.py

app.py

@@ -1,10 +1,11 @@
-from transformers import LongformerTokenizerFast, LongformerForSequenceClassification
+import transformers
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM
 import torch
 import gradio as gr
 import requests
+import bs4
 from bs4 import BeautifulSoup
 
-
 def get_text_from_url(url):
 
     headers = {
@@ -20,14 +21,26 @@ def get_text_from_url(url):
     else:
 
         print("Error al obtener la página:", response.status_code)
-        return None
+        return 'error'
 
 classification_model_checkpoint = 'FrancoMartino/privacyPolicies_classification'
-classification_tokenizer = LongformerTokenizerFast.from_pretrained(classification_model_checkpoint)
-classification_model = LongformerForSequenceClassification.from_pretrained(classification_model_checkpoint)
+classification_tokenizer = AutoTokenizer.from_pretrained("FrancoMartino/privacyPolicies_classification")
+classification_model = AutoModelForSequenceClassification.from_pretrained("FrancoMartino/privacyPolicies_classification")
+
+summarization_model_checkpoint = "facebook/bart-large-cnn"
+summarization_tokenizer = AutoTokenizer.from_pretrained(summarization_model_checkpoint)
+summarization_model = AutoModelForSeq2SeqLM.from_pretrained(summarization_model_checkpoint)
 
 def predict(url):
     text = get_text_from_url(url)
+    if text == 'error':
+      return {'ERROR': 'Error with the url'}
+    if len(classification_tokenizer.tokenize(text)) > 4096:
+      print('long')
+      inputs = summarization_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+      with torch.no_grad():
+          summary_ids = summarization_model.generate(inputs['input_ids'], max_length=4096)
+      text = summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     inputs = classification_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=4096)
     with torch.no_grad():
         logits = classification_model(**inputs).logits