feat: cpu support using cuda condition branch

app.py

@@ -2,6 +2,7 @@ import os
 
 import streamlit as st
 from pymilvus import MilvusClient
+import torch
 
 from model import encode_dpr_question, get_dpr_encoder
 from model import summarize_text, get_summarizer
@@ -56,6 +57,7 @@ st.markdown(styl, unsafe_allow_html=True)
 question = st.text_area("Text to summarize", INITIAL, height=400)
 
 
+@torch.inference_mode()
 def main(question: str):
   if question in st.session_state:
     print("Cache hit!")

model.py

@@ -7,15 +7,21 @@ from transformers import QuestionAnsweringPipeline
 from transformers import AutoTokenizer, PegasusXForConditionalGeneration, PegasusTokenizerFast
 import torch
 
+cuda = torch.cuda.is_available()
 max_answer_len = 8
 logging.set_verbosity_error()
 
 
+@torch.inference_mode()
 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):
+                     return_tensors='pt', truncation=True)
+  if cuda:
+    inputs = inputs.to(0)
+    with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
+      summary_ids = model.generate(inputs["input_ids"])
+  else:
     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))
@@ -24,14 +30,13 @@ def summarize_text(tokenizer: PegasusTokenizerFast, model: PegasusXForConditiona
 
 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 = PegasusXForConditionalGeneration.from_pretrained(model_id)
+  if cuda:
+    model = model.to(0)
   model = torch.compile(model)
   return tokenizer, model
 
 
-# OpenAI reader
-
-
 class AnswerInfo(TypedDict):
   score: float
   start: int
@@ -42,10 +47,16 @@ class AnswerInfo(TypedDict):
 @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):
+  if cuda:
+    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)
+  else:
     pipeline = QuestionAnsweringPipeline(
-        model=model, tokenizer=tokenizer, device='cuda', max_answer_len=max_answer_len)
-    answer_infos: List[AnswerInfo] = pipeline(
+        model=model, tokenizer=tokenizer, device='cpu', max_answer_len=max_answer_len)
+    answer_infos = pipeline(
         question=questions, context=ctxs)
   for answer_info in answer_infos:
     answer_info['answer'] = sub(r'[.\(\)"\',]', '', answer_info['answer'])
@@ -54,10 +65,13 @@ def ask_reader(tokenizer: AutoTokenizer, model: AutoModelForQuestionAnswering,
 
 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)
+  model = DPRReader.from_pretrained(model_id)
+  if cuda:
+    model = model.to(0)
   return tokenizer, model
 
 
+@torch.inference_mode()
 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
@@ -67,9 +81,13 @@ def encode_dpr_question(tokenizer: DPRQuestionEncoderTokenizer, model: DPRQuesti
       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
+                         padding=True, truncation=True)
+  if cuda:
+    batch_dict = batch_dict.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
+  else:
+    embeddings = model(**batch_dict).pooler_output
   return embeddings
 
 
@@ -82,5 +100,7 @@ def get_dpr_encoder(model_id="facebook/dpr-question_encoder-single-nq-base") ->
       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)
+  model = DPRQuestionEncoder.from_pretrained(model_id)
+  if cuda:
+    model = model.to(0)
   return tokenizer, model