app.py

# 使用gradio开发QA的可视化demo

import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering, BigBirdForQuestionAnswering, BigBirdConfig, PreTrainedModel, BigBirdTokenizer
import torch
from torch import nn
from transformers.models.big_bird.modeling_big_bird import BigBirdOutput, BigBirdIntermediate

class BigBirdNullHead(nn.Module):
    """Head for question answering tasks."""

    def __init__(self, config):
        super().__init__()
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.intermediate = BigBirdIntermediate(config)
        self.output = BigBirdOutput(config)
        self.qa_outputs = nn.Linear(config.hidden_size, 2)

    def forward(self, encoder_output):
        hidden_states = self.dropout(encoder_output)
        hidden_states = self.intermediate(hidden_states)
        hidden_states = self.output(hidden_states, encoder_output)
        logits = self.qa_outputs(hidden_states)
        return logits


model_path = './checkpoint-epoch-best'

class BigBirdForQuestionAnsweringWithNull(PreTrainedModel):
    def __init__(self, config, model_id):
        super().__init__(config)
        self.bertqa = BigBirdForQuestionAnswering.from_pretrained(model_id,
            config=self.config, add_pooling_layer=True)
        self.null_classifier = BigBirdNullHead(self.bertqa.config)
        self.contrastive_mlp = nn.Sequential(
            nn.Linear(self.bertqa.config.hidden_size, self.bertqa.config.hidden_size),
        )

    def forward(self, **kwargs):
        if self.training:
            null_labels = kwargs['is_impossible']
            del kwargs['is_impossible']
            outputs = self.bertqa(**kwargs)
            pooler_output = outputs.pooler_output
            null_logits = self.null_classifier(pooler_output)
            loss_fct = nn.CrossEntropyLoss()
            null_loss = loss_fct(null_logits, null_labels)


            outputs.loss = outputs.loss + null_loss 

            return outputs.to_tuple()
        else:
            outputs = self.bertqa(**kwargs)
            pooler_output = outputs.pooler_output
            null_logits = self.null_classifier(pooler_output)

            return (outputs.start_logits, outputs.end_logits, null_logits)
model_id = 'vasudevgupta/bigbird-roberta-natural-questions'
config = BigBirdConfig.from_pretrained(model_id)
model = BigBirdForQuestionAnsweringWithNull(config, model_id)
# model.to('cuda')
model.eval()

model.load_state_dict(torch.load(model_path+'/pytorch_model.bin',map_location=torch.device('cpu'))) # map_location是指定加载到哪个设备

tokenizer = BigBirdTokenizer.from_pretrained(model_path)

def main(question, context):
    # 编码输入
    text = question + " [SEP] " + context
    inputs = tokenizer(text, max_length=4096, truncation=True, return_tensors="pt")
    print(question)
    # inputs.to('cuda')
    # 预测答案
    outputs = model(**inputs)
    start_scores = outputs[0]
    end_scores = outputs[1]
    null_scores = outputs[2]
    # 解码答案
    is_impossible = null_scores.argmax().item()
    if is_impossible:
        print('No Answer')
        return "No Answer"
    else:
        answer_start = torch.argmax(start_scores)
        answer_end = torch.argmax(end_scores) + 1
        answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs['input_ids'][0][answer_start:answer_end]))
        print(answer)
        return answer

  

with gr.Blocks() as demo:

  gr.Markdown("""# TimeQA base!""")
  gr.Markdown("""#### 数据集链接 [TimeQA](https://czq1999.github.io/HTML/research/doc_can.html)""")
  with gr.Row():
    with gr.Column():
    #   options = gr.inputs.Radio(["vasudevgupta/bigbird-roberta-natural-questions", "vasudevgupta/bigbird-roberta-natural-questions"], label="Model")
      text1 = gr.Textbox(
            label="Question",
            lines=1,
            value="Which team did the player Sean Bowers belong to from 1996 to 1997?",
        )
      text2 = gr.Textbox(
            label="Context",
            lines=6,
            value="Sean Bowers Sean Patrick Bowers ( born August 12 , 1968 ) is a retired U.S . soccer defender . He spent most of his career playing indoor soccer , earning the 1992 NPSL Rookie of the Year , four-time Defender of the Year and a six-time first team All Star in four different indoor leagues . He also played four seasons with the Kansas City Wizards in Major League Soccer . He currently plays for the San Diego Sockers of the Professional Arena Soccer League . High school and college . Bowers grew up in San Diego , California where he played his youth soccer with the San Diego Nomads . He also played goalkeeper at Mira Mesa High School where he was two-time All City . Following high school , he attended Miramar College for two years before finishing college at Quincy University in Quincy , Illinois . He played soccer at both Miramar and Quincy . In 1990 , Bowers led Quincy in scoring , being named the team MVP . He graduated in 1991 with a bachelors degree in political science . Player . In 1991 , Bowers signed with the San Diego Sockers of the Major Indoor Soccer League ( MISL ) . While he played in the Sockers season opener , at some point he moved to the Detroit Rockers of the National Professional Soccer League ( NPSL ) . He was selected as the NPSL Rookie of the Year . He remained with the Rockers until 1995 , winning 1994 and 1995 Defender of the Year honors . In addition to playing with the Rockers in 1993 , Bowers also played a single season with the Sacramento Knights of the Continental Indoor Soccer League ( CISL ) , a summer indoor league . He was named CISL Defender of the Year . In 1995 and 1996 , he returned to the CISL , this time with the Anaheim Splash . The Kansas City Wiz of the outdoor Major League Soccer selected Bowers in the third round ( 26th overall ) in the 1996 MLS Inaugural Player Draft . He played four seasons with the Wizards , serving as the team captain in 1999 . However , he continued to play indoor soccer , returning to the Detroit Rockers in 1997 . He played with the Rockers until 2001 , when the team folded . The Baltimore Blast selected Bowers in the dispersal draft . That summer , the NPSL also gave way to the new Major Indoor Soccer League ( MISL ) . Bowers spent two season with the Blast , winning the 2003 league championship with them . July 16 , 2003 , the Blast traded Bowers to the San Diego Sockers in exchange for Carlos Farias and the first and second picks in the 2003 Amateur Draft . He played thirty-three games with the Sockers in 2003–2005 , but only five games the next season before the Sockers folded . Rather than pursuing his playing career further , Bowers elected to retire , settle in his hometown and pursue a coaching career . In 2009 , he joined the San Diego Sockers of the Professional Arena Soccer League . Futsal . Beginning in 1996 , Bowers became a regular part of the U.S . national futsal team . Two years later , the team took third place in the Futsal Mundialito . He earned a total of thirty-six caps , scoring five goals , with the U.S . from 1996 to 2004 . Coaching . In 1999 , Bowers became an assistant coach with William Jewell College womens soccer team . This was the start of his coaching career . Bowers began coaching on the professional level with the Detroit Rockers in 2000–2001 . That season , he served as an assistant coach in addition to playing . In 2001 , he became the head coach with the Detroit Rocker Hawks , winners of the 2001 USASA Womens U.S . Open Cup . In 2002 , he became the interim head coach of the Baltimore Blast . At the time , the Blast were at the bottom of the standings . Bowers took them into the playoffs and the Blast ended up winning the MISL championship that season . In 2003 , the Blast traded Bowers to the San Diego Sockers . When Bowers arrived in southern California , he became an assistant coach with the California Baptist University womens team from 2003 to 2005 . He was elevated to the position of head coach in 2006 , taking them to an end of season ranking of ninth in the NAIA . He also was named the head coach of the San Diego SeaLions of the Womens Premier Soccer League ( WPSL ) in 2006 . That year the SeaLions went to the WPSL semifinals . On March 27 , 2007 , Miramar College announced the creation of a womens soccer team with Bowers as its coach . Honors . Rookie of the Year Defender of the Year First Team All Star External links . - Indoor soccer honors - Southwest Soccer Club profile - San Diego Sockers bio",
        )
      output = gr.Textbox(
          label="Answer",
      )
      b1 = gr.Button("Ask Question!")
      b1.click(main, inputs=[text1, text2], outputs=output)
#   gr.Markdown("""#### powered by [Tassle](https://bit.ly/3LXMklV)""")


if __name__ == "__main__":

    demo.launch()