app.py

import gradio as gr
import os
import torch
import transformers

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

HF_TOKEN_DOWNLOAD = os.environ.get('HF_TOKEN_DOWNLOAD')

class Processor:
    def __init__(self, model):
        self.tokenizer = transformers.AutoTokenizer.from_pretrained(model, use_auth_token=HF_TOKEN_DOWNLOAD)
        self.model = transformers.AutoModelForSeq2SeqLM.from_pretrained(model, use_auth_token=HF_TOKEN_DOWNLOAD, low_cpu_mem_usage=True, device_map='auto', torch_dtype='auto', offload_folder='offload')
        self.model.eval()

    def parse_choices(self, s):
        '''
        s: serialized_choices '(A) ... (B) ... (C) ...'
        '''
        choices = []
        key = 'A' if s.find('(A)') != -1 else 'a'
        while True:
            pos = s.find(f'({chr(ord(key) + 1)})')
            if pos == -1:
                break
            choice = s[3:pos]
            s = s[pos:]
            choice = choice.strip(' ')
            choices.append(choice)
            key = chr(ord(key) + 1)
        choice = s[3:]
        choice = choice.strip(' ')
        choices.append(choice)
        return choices

    def run(self, question, max_question_len, max_knowledge_len, max_answer_len, m, top_p):
        choices = self.parse_choices(question.split('\\n')[1].strip(' '))
        choices = [chr(ord('A') + i) for i, choice in enumerate(choices)]
        choices_ids = self.tokenizer(choices, return_tensors='pt', padding='max_length', truncation='longest_first', max_length=max_answer_len).input_ids.to(device) # (C, AL)

        prompt = question + ' \\n Knowledge: '
        prompt_tok = self.tokenizer(prompt, return_tensors='pt', padding='max_length', truncation='longest_first', max_length=max_question_len).to(device) # (1, QL)
        knowledges_ids = self.model.generate(
            input_ids=prompt_tok.input_ids,
            attention_mask=prompt_tok.attention_mask,
            max_length=max_knowledge_len + 1,
            min_length=3,
            do_sample=True,
            num_return_sequences=m,
            top_p=top_p,
        ) # (K, KL); begins with 0 ([BOS]); ends with 1 ([EOS])
        knowledges_ids = knowledges_ids[:, 1:].contiguous() # no beginning; ends with 1 ([EOS])
        knowledges = self.tokenizer.batch_decode(knowledges_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
        knowledges = list(set(knowledges))
        knowledges = [''] + knowledges

        prompts = [question + (f' \\n Knowledge: {knowledge} \\n Answer: ' if knowledge != '' else ' \\n Answer:') for knowledge in knowledges]
        prompts_tok = self.tokenizer(prompts, return_tensors='pt', padding='max_length', truncation='longest_first', max_length=max_question_len + max_knowledge_len).to(device) # (1+K, QL+KL)
        output = self.model(
            input_ids=prompts_tok.input_ids,
            attention_mask=prompts_tok.attention_mask,
            labels=choices_ids[0].unsqueeze(0).repeat(len(knowledges), 1),
        )
        logitsss = output.logits # (1+K, AL, V)
        logitss = logitsss[:, 0, :] # (1+K, V)
        choice_ids = choices_ids[:, 0] # (C)
        answer_logitss = logitss.gather(dim=1, index=choice_ids.unsqueeze(0).expand(len(knowledges), -1)) # (1+K, C)
        answer_probss = answer_logitss.softmax(dim=1) # (1+K, C)

        # Ensemble
        knowless_pred = answer_probss[0, :].argmax(dim=0).item()
        knowless_pred = choices[knowless_pred]

        answer_probs = answer_probss.max(dim=0).values # (C)
        knowful_pred = answer_probs.argmax(dim=0).item()
        knowful_pred = choices[knowful_pred]
        selected_knowledge_ix = answer_probss.max(dim=1).values.argmax(dim=0).item()
        selected_knowledge = knowledges[selected_knowledge_ix]

        return {
            'question': question,
            'knowledges': knowledges,
            'knowless_pred': knowless_pred,
            'knowful_pred': knowful_pred,
            'selected_knowledge': selected_knowledge,
        }

MODELS = [
    'liujch1998/crystal-large',
    # 'liujch1998/crystal-3b',
    # 'liujch1998/crystal-11b',
]
processor_by_model = {}
for model in MODELS:
    processor_by_model[model] = Processor(model)

def predict(question, model, max_question_len, max_knowledge_len, max_answer_len, m, top_p):
    result = processor_by_model[model].run(question, max_question_len, max_knowledge_len, max_answer_len, m, top_p)
    return result['knowless_pred'], result['knowful_pred'], '\n'.join(result['knowledges']), result['selected_knowledge']

examples = [
    'If the mass of an object gets bigger what will happen to the amount of matter contained within it? \\n (A) gets bigger (B) gets smaller',
    'What would vinyl be an odd thing to replace? \\n (A) pants (B) record albums (C) record store (D) cheese (E) wallpaper',
    'Some pelycosaurs gave rise to reptile ancestral to \\n (A) lamphreys (B) angiosperm (C) mammals (D) paramecium (E) animals (F) protozoa (G) arachnids (H) backbones',
    'Sydney rubbed Addison’s head because she had a horrible headache. What will happen to Sydney? \\n (A) drift to sleep (B) receive thanks (C) be reprimanded',
    'Adam always spent all of the free time watching Tv unlike Hunter who volunteered, due to _ being lazy. \\n (A) Adam (B) Hunter',
    'Causes bad breath and frightens blood-suckers \\n (A) tuna (B) iron (C) trash (D) garlic (E) pubs',
]

input_question = gr.Dropdown(choices=examples, label='Question:',
    info='A multiple-choice commonsense question. Please follow the UnifiedQA input format: "{question} \\n (A) ... (B) ... (C) ..."',
)
input_model = gr.Dropdown(label='Model:', value=MODELS[0], choices=MODELS)
input_max_question_len = gr.Number(label='Max number of tokens in question:', value=256, precision=0)
input_max_knowledge_len = gr.Number(label='Max number of tokens in knowledge:', value=32, precision=0)
input_max_answer_len = gr.Number(label='Max number of tokens in answer:', value=2, precision=0)
input_m = gr.Slider(label='Number of generated knowledges:', value=10, mininum=1, maximum=20, step=1,
    info='The actual number of generated knowledges may be less than this number due to possible duplicates.',
)
input_top_p = gr.Slider(label='top_p for knowledge generation:', value=0.5, mininum=0.0, maximum=1.0, step=0.05)
output_knowless_answer = gr.Textbox(label='QA model answer without knowledge:', interactive=False)
output_knowful_answer = gr.Textbox(label='QA model answer with knowledge:', interactive=False)
output_all_knowledges = gr.Textbox(label='All generated knowledges:', interactive=False)
output_selected_knowledge = gr.Textbox(label='Knowledge selected to make the prediction:', interactive=False)

description = '''This is a demo for the paper, [*Crystal: Introspective Reasoners Reinforced with Self-Feedback*](https://arxiv.org/abs/2310.04921), presented at EMNLP 2023. [[Code](https://github.com/liujch1998/crystal)] [[Model](https://huggingface.co/liujch1998/crystal-11b)] This demo is made & maintained by [Jiacheng (Gary) Liu](https://liujch1998.github.io).
Crystal is an introspective reasoning model that answers commonsense questions by first generating knowledge and then use knowledge-grounded reasoning to reach a final prediction. To try this model, select an example question, or write your own commonsense question in the suggested format.'''

gr.Interface(
    fn=predict,
    inputs=[input_question, input_model, input_max_question_len, input_max_knowledge_len, input_max_answer_len, input_m, input_top_p],
    outputs=[output_knowless_answer, output_knowful_answer, output_all_knowledges, output_selected_knowledge],
    title="Crystal Demo",
    description=description,
).launch()