use ms2 for summarization

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, LEDForConditionalGeneration
 import requests
 from bs4 import BeautifulSoup
 import nltk
@@ -149,10 +149,11 @@ def init_models():
     reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-2-v2', device=device)
     # queryexp_tokenizer = AutoTokenizer.from_pretrained("doc2query/all-with_prefix-t5-base-v1")
     # queryexp_model = AutoModelWithLMHead.from_pretrained("doc2query/all-with_prefix-t5-base-v1")
-    summarizer = pipeline("summarization")
-    return question_answerer, reranker, stop, device, summarizer
+    summ_tok = AutoTokenizer.from_pretrained('allenai/led-base-16384-ms2')
+    summ_mdl = LEDForConditionalGeneration.from_pretrained('allenai/led-base-16384-ms2')
+    return question_answerer, reranker, stop, device, summ_mdl, summ_tok
 
-qa_model, reranker, stop, device, summarizer = init_models() # queryexp_model, queryexp_tokenizer
+qa_model, reranker, stop, device, summ_mdl, summ_tok = init_models() # queryexp_model, queryexp_tokenizer
 
 
 def clean_query(query, strict=True, clean=True):
@@ -270,15 +271,71 @@ def matched_context(start_i, end_i, contexts_string, seperator='---'):
             return contexts_string[doc_starts[i]:doc_starts[i+1]].replace(seperator, '')
     return None
 
+def process_document(documents, tokenizer, docsep_token_id, pad_token_id, device=device):
+    input_ids_all=[]
+    for data in documents:
+        all_docs = data.split("|||||")
+        for i, doc in enumerate(all_docs):
+            doc = doc.replace("\n", " ")
+            doc = " ".join(doc.split())
+            all_docs[i] = doc
+
+        #### concat with global attention on doc-sep
+        input_ids = []
+        for doc in all_docs:
+            input_ids.extend(
+                tokenizer.encode(
+                    doc,
+                    truncation=True,
+                    max_length=4096 // len(all_docs),
+                )[1:-1]
+            )
+            input_ids.append(docsep_token_id)
+        input_ids = (
+            [tokenizer.bos_token_id]
+            + input_ids
+            + [tokenizer.eos_token_id]
+        )
+        input_ids_all.append(torch.tensor(input_ids))
+    input_ids = torch.nn.utils.rnn.pad_sequence(
+        input_ids_all, batch_first=True, padding_value=pad_token_id
+    )
+    return input_ids
+
+
+def batch_process(batch, model, tokenizer, docsep_token_id, pad_token_id, device=device):
+    input_ids=process_document(batch['document'], tokenizer, docsep_token_id, pad_token_id)
+    # get the input ids and attention masks together
+    global_attention_mask = torch.zeros_like(input_ids).to(device)
+    input_ids = input_ids.to(device)
+    # put global attention on <s> token
+
+    global_attention_mask[:, 0] = 1
+    global_attention_mask[input_ids == docsep_token_id] = 1
+    generated_ids = model.generate(
+        input_ids=input_ids,
+        global_attention_mask=global_attention_mask,
+        use_cache=True,
+        max_length=1024,
+        num_beams=5,
+    )
+    generated_str = tokenizer.batch_decode(
+            generated_ids.tolist(), skip_special_tokens=True
+        )
+    result={}
+    result['generated_summaries'] = generated_str
+    return result
+
 
 def gen_summary(query, sorted_result):
-    doc_sep = '\n'
-    summary = summarizer(f'{query} '.join([f'{doc_sep}'.join(r['texts']) + r['context'] for r in sorted_result]))[0]['summary_text']
+    pad_token_id = summ_tok.pad_token_id
+    docsep_token_id = summ_tok.convert_tokens_to_ids("</s>")
+    out = batch_process({ 'document': [f'||||'.join([f'{query} '.join(r['texts']) + r['context'] for r in sorted_result])]}, summ_mdl, summ_tok, docsep_token_id, pad_token_id)
     st.markdown(f"""
     <div class="container-fluid">
     <div class="row align-items-start">
             <div  class="col-md-12 col-sm-12">
-        <strong>Answer:</strong> {summary}
+        <strong>Answer:</strong> {out['generated_summaries'][0]}
         </div>
     </div>
     </div>