Spaces:
Sleeping
Sleeping
File size: 3,648 Bytes
75148a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
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
|