remove summarization

app.py

@@ -78,7 +78,6 @@ def search(term, limit=10, clean=True, strict=True, all_mode=True, abstracts=Tru
         except:
             pass
 
-
     return (
         contexts,
         docs
@@ -149,11 +148,9 @@ 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")
-    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
+    return question_answerer, reranker, stop, device
 
-qa_model, reranker, stop, device, summ_mdl, summ_tok = init_models() # queryexp_model, queryexp_tokenizer
+qa_model, reranker, stop, device = init_models() # queryexp_model, queryexp_tokenizer
 
 
 def clean_query(query, strict=True, clean=True):
@@ -214,9 +211,6 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 with st.expander("Settings (strictness, context limit, top hits)"):
-    use_mds = st.radio(
-        "Use multi-document summarization to summarize answer?",
-        ('yes', 'no'))
     support_all = st.radio(
         "Use abstracts and titles as a ranking signal (if the words are matched in the abstract then the document is more relevant)?",
         ('yes', 'no'))
@@ -271,77 +265,6 @@ 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):
-    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> {out['generated_summaries'][0]}
-        </div>
-    </div>
-    </div>
-    """, unsafe_allow_html=True)
-    st.markdown("<br /><br /><h5>Sources:</h5>", unsafe_allow_html=True)
-
 
 def run_query(query):
 #     if use_query_exp == 'yes':
@@ -395,7 +318,7 @@ def run_query(query):
         context = '\n---'.join(contexts[:context_limit])
 
     results = []
-    model_results = qa_model(question=query, context=context, top_k=10)
+    model_results = qa_model(question=query, context=query+'---'+context, top_k=10)
     for result in model_results:
         matched = matched_context(result['start'], result['end'], context)
         support = find_source(result['answer'], orig_docs, matched)
@@ -423,9 +346,6 @@ def run_query(query):
         sorted_result
     ))
 
-    if use_mds == 'yes':
-        gen_summary(query, sorted_result)
-
     for r in sorted_result:
         ctx = remove_html(r["context"])
         for answer in r['texts']: