DevashishBhake/Question_Generation

1import torch2from transformers import T5Tokenizer, T5ForConditionalGeneration,T5Model3import gradio as gr4 5def get_questions(paragraph, tokenizer, model, device):6    bt_levels = ['Remember', 'Understand', 'Apply', 'Analyse', 'Evaluate', 'Create']7    questions_dict = {}8    for bt_level in bt_levels:9        input_text = f'{bt_level}: {paragraph} {tokenizer.eos_token}'10        input_ids = tokenizer.encode(input_text, max_length=512, padding='max_length', truncation=True, return_tensors='pt').to(device)11        model.eval()12        generated_ids = model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True).to(device)13        output_text = tokenizer.decode(generated_ids.squeeze(), skip_special_tokens=True).lstrip('\n')14        output_text = output_text.split(' ', 1)[1]15        questions_dict.update({bt_level: output_text})16        # print(f'{bt_level} level question: {output_text}')17    return questions_dict18        19 20def main(paragraph):21    model = T5ForConditionalGeneration.from_pretrained('./save_model')22    tokenizer = T5Tokenizer.from_pretrained('./save_model')23    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')24    model.to(device)25    output = get_questions(paragraph, tokenizer, model, device)26    return output27 28gr.Interface(29    fn=main, 30    inputs="textbox",31    outputs="textbox",32    live=True).launch()