nickmuchi commited on
Commit
93c8623
1 Parent(s): d8b4368

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -6
app.py CHANGED
@@ -325,14 +325,16 @@ def get_spacy():
325
 
326
  @st.experimental_singleton(suppress_st_warning=True)
327
  def facebook_model():
328
-
329
- summarizer = pipeline('summarization',model='facebook/bart-large-cnn')
 
330
  return summarizer
331
 
332
  @st.experimental_singleton(suppress_st_warning=True)
333
  def schleifer_model():
334
-
335
- summarizer = pipeline('summarization',model='sshleifer/distilbart-cnn-12-6')
 
336
  return summarizer
337
 
338
  #@st.experimental_singleton(suppress_st_warning=True)
@@ -447,8 +449,8 @@ if summarize:
447
  text="Loading Facebook-Bart Model and Extracting summary. This might take a few seconds depending on the length of your text..."
448
  ):
449
  summarizer_model = facebook_model()
450
- summarized_text = summarizer_model(text_to_summarize, max_length=max_len, min_length=min_len)
451
- summarized_text = ' '.join([summ['summary_text'] for summ in summarized_text])
452
 
453
  elif model_type == "Sshleifer-DistilBart":
454
  if url_text:
 
325
 
326
  @st.experimental_singleton(suppress_st_warning=True)
327
  def facebook_model():
328
+ model_name = 'facebook/bart-large-cnn'
329
+ summarizer = pipeline('summarization',model=model_name,tokenizer=model_name,
330
+ device=0 if torch.cuda.is_available() else -1)
331
  return summarizer
332
 
333
  @st.experimental_singleton(suppress_st_warning=True)
334
  def schleifer_model():
335
+ model_name = 'sshleifer/distilbart-cnn-12-6'
336
+ summarizer = pipeline('summarization',model=model_name, tokenizer=model_name,
337
+ device=0 if torch.cuda.is_available() else -1)
338
  return summarizer
339
 
340
  #@st.experimental_singleton(suppress_st_warning=True)
 
449
  text="Loading Facebook-Bart Model and Extracting summary. This might take a few seconds depending on the length of your text..."
450
  ):
451
  summarizer_model = facebook_model()
452
+ summarized_text = summarizer_model(text_to_summarize.strip().replace("\n", " "), max_length=max_len, min_length=min_len,clean_up_tokenization_spaces=True)
453
+ summarized_text = ' '.join([summ['summary_text'].replace("<n>", " ") for summ in summarized_text])
454
 
455
  elif model_type == "Sshleifer-DistilBart":
456
  if url_text: