Update app.py

app.py

@@ -12,17 +12,21 @@ import pdfkit
 import yake
 from summarizer import Summarizer,TransformerSummarizer
 from transformers import pipelines
-#nltk.download('punkt')
-
+nltk.download('punkt')
 from transformers import AutoTokenizer, AutoModelForPreTraining
+# model_name = 'distilbert-base-uncased'
 model_name = 'nlpaueb/legal-bert-base-uncased'
+#model_name = 'laxya007/gpt2_legal'
+# model_name = 'facebook/bart-large-cnn'
 
+# The setup of huggingface.co
+custom_config = AutoConfig.from_pretrained(model_name)
+custom_config.output_hidden_states=True
+custom_tokenizer = AutoTokenizer.from_pretrained(model_name)
+custom_model = AutoModel.from_pretrained(model_name, config=custom_config)
+bert_legal_model = Summarizer(custom_model=custom_model, custom_tokenizer=custom_tokenizer)
+print('Using model {}\n'.format(model_name))
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-model = AutoModelForPreTraining.from_pretrained(model_name)
-bert_legal_model = Summarizer(custom_model=model, custom_tokenizer=tokenizer)
- 
 
 
 def get_response(input_text):