feat: huggingface space pipeline with resrer model

app.py

@@ -1,2 +1,97 @@
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+import os
+
+import streamlit as st
+from pymilvus import MilvusClient
+
+from model import encode_dpr_question, get_dpr_encoder
+from model import summarize_text, get_summarizer
+from model import ask_reader, get_reader
+
+
+TITLE = 'ReSRer: Retriever-Summarizer-Reader'
+INITIAL = "What is the population of NYC"
+
+st.set_page_config(page_title=TITLE)
+st.header(TITLE)
+st.markdown('''
+### Ask short-answer question that can be find in Wikipedia data.
+''', unsafe_allow_html=True)
+
+
+@st.cache_resource
+def load_models():
+  models = {}
+  models['encoder'] = get_dpr_encoder()
+  models['summarizer'] = get_summarizer()
+  models['reader'] = get_reader()
+  return models
+
+
+@st.cache_resource
+def load_client():
+  client = MilvusClient(user='resrer', password=os.env['MILVUS_PW'],
+                        uri=f"http://{os.env['MILVUS_HOST']}:19530", db_name='psgs_w100')
+  return client
+
+
+client = load_client()
+models = load_models()
+
+styl = """
+<style>
+    .StatusWidget-enter-done{
+      position: fixed;
+      left: 50%;
+      top: 50%;
+      transform: translate(-50%, -50%);
+    }
+    .StatusWidget-enter-done button{
+      display: none;
+    }
+</style>
+"""
+st.markdown(styl, unsafe_allow_html=True)
+
+
+question = st.text_area("Text to summarize", INITIAL, height=400)
+
+
+def main(question: str):
+  if question in st.session_state:
+    print("Cache hit!")
+    ctx, summary, answer = st.session_state[question]
+  else:
+    print(f"Input: {question}")
+    # Embedding
+    question_vectors = encode_dpr_question(
+        models['encoder'][0], models['encoder'][1], [question])
+    query_vector = question_vectors.detach().cpu().numpy().tolist()[0]
+
+    # Retriever
+    results = client.search(collection_name='dpr_nq', data=[
+        query_vector], limit=10, output_fields=['title', 'text'])
+    texts = [result['entity']['text'] for result in results[0]]
+    ctx = '\n'.join(texts)
+
+    # Reader
+    summary = summarize_text(models['summarizer'][0],
+                             models['summarizer'][1], [summary])
+    answers = ask_reader(models['reader'][0],
+                         models['reader'][1], [question], [ctx])
+    answer = answers[0]['answer']
+    print(f"\nAnswer: {answer}")
+
+    st.session_state[question] = (ctx, summary, answer)
+
+  # Summary
+  st.markdown(answer)
+  st.write("## Summary")
+  st.markdown(
+      f"<h6 style='padding: 0'>{summary}</h6><hr style='margin: 1em 0px'>", unsafe_allow_html=True)
+  st.markdown(ctx)
+
+  st.write(f"{question}", unsafe_allow_html=True)
+
+
+if question:
+  main(question)

model.py

@@ -0,0 +1,86 @@
+from typing import List, Tuple, TypedDict
+from re import sub
+
+from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer, logging
+from transformers import AutoModelForQuestionAnswering, DPRReaderTokenizer, DPRReader
+from transformers import QuestionAnsweringPipeline
+from transformers import AutoTokenizer, PegasusXForConditionalGeneration, PegasusTokenizerFast
+import torch
+
+max_answer_len = 8
+logging.set_verbosity_error()
+
+
+def summarize_text(tokenizer: PegasusTokenizerFast, model: PegasusXForConditionalGeneration,
+                   input_texts: List[str]):
+  inputs = tokenizer(input_texts, padding=True,
+                     return_tensors='pt', truncation=True).to(1)
+  with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
+    summary_ids = model.generate(inputs["input_ids"])
+  summaries = tokenizer.batch_decode(summary_ids, skip_special_tokens=True,
+                                     clean_up_tokenization_spaces=False, batch_size=len(input_texts))
+  return summaries
+
+
+def get_summarizer(model_id="seonglae/resrer") -> Tuple[PegasusTokenizerFast, PegasusXForConditionalGeneration]:
+  tokenizer = PegasusTokenizerFast.from_pretrained(model_id)
+  model = PegasusXForConditionalGeneration.from_pretrained(model_id).to(1)
+  model = torch.compile(model)
+  return tokenizer, model
+
+
+# OpenAI reader
+
+
+class AnswerInfo(TypedDict):
+  score: float
+  start: int
+  end: int
+  answer: str
+
+
+@torch.inference_mode()
+def ask_reader(tokenizer: AutoTokenizer, model: AutoModelForQuestionAnswering,
+               questions: List[str], ctxs: List[str]) -> List[AnswerInfo]:
+  with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
+    pipeline = QuestionAnsweringPipeline(
+        model=model, tokenizer=tokenizer, device='cuda', max_answer_len=max_answer_len)
+    answer_infos: List[AnswerInfo] = pipeline(
+        question=questions, context=ctxs)
+  for answer_info in answer_infos:
+    answer_info['answer'] = sub(r'[.\(\)"\',]', '', answer_info['answer'])
+  return answer_infos
+
+
+def get_reader(model_id="mrm8488/longformer-base-4096-finetuned-squadv2"):
+  tokenizer = DPRReaderTokenizer.from_pretrained(model_id)
+  model = DPRReader.from_pretrained(model_id).to(0)
+  return tokenizer, model
+
+
+def encode_dpr_question(tokenizer: DPRQuestionEncoderTokenizer, model: DPRQuestionEncoder, questions: List[str]) -> torch.FloatTensor:
+  """Encode a question using DPR question encoder.
+  https://huggingface.co/docs/transformers/model_doc/dpr#transformers.DPRQuestionEncoder
+
+  Args:
+      question (str): question string to encode
+      model_id (str, optional): Default for NQ or "facebook/dpr-question_encoder-multiset-base
+  """
+  batch_dict = tokenizer(questions, return_tensors="pt",
+                         padding=True, truncation=True,).to(0)
+  with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
+    embeddings: torch.FloatTensor = model(**batch_dict).pooler_output
+  return embeddings
+
+
+def get_dpr_encoder(model_id="facebook/dpr-question_encoder-single-nq-base") -> Tuple[DPRQuestionEncoder, DPRQuestionEncoderTokenizer]:
+  """Encode a question using DPR question encoder.
+  https://huggingface.co/docs/transformers/model_doc/dpr#transformers.DPRQuestionEncoder
+
+  Args:
+      question (str): question string to encode
+      model_id (str, optional): Default for NQ or "facebook/dpr-question_encoder-multiset-base
+  """
+  tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(model_id)
+  model = DPRQuestionEncoder.from_pretrained(model_id).to(0)
+  return tokenizer, model

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+torch
+pymilvus