wasimmadha/entity-extraction

1if __name__ == '__main__':    2    inputs = ['gbjjhbdjhbdgjhdbfjhsdkjrkjf', 'fdjhbjhsbd']3    from transformers import AutoTokenizer4    from model import CustomModel5    import torch6    from configuration import CFG7    from dataset import SingleInputDataset8    from torch.utils.data import DataLoader9    from utils import inference_fn, get_char_probs, get_results, get_text10    import numpy as np11    import gradio as gr12    import os13 14    device = torch.device('cpu')15    config_path = os.path.join('models_file', 'config.pth')16    model_path  = os.path.join('models_file', 'microsoft-deberta-base_0.9449373420387531_8_best.pth')17    tokenizer = AutoTokenizer.from_pretrained('models_file/tokenizer')18    model = CustomModel(CFG, config_path=config_path, pretrained=False)19    state = torch.load(model_path,20                        map_location=torch.device('cpu'))21    model.load_state_dict(state['model'])22 23    def get_answer(context, feature):24 25        ## Input to the model using patient-history and feature-text26        inputs_single = tokenizer(context, feature, 27                                    add_special_tokens=True,28                                    max_length=CFG.max_len,29                                    padding="max_length",30                                    return_offsets_mapping=False)31 32        for k, v in inputs_single.items():33            inputs_single[k] = torch.tensor(v, dtype=torch.long)34 35        # Create a new dataset containing only the input sample36        single_input_dataset = SingleInputDataset(inputs_single)37        # Create a DataLoader for the new dataset38        single_input_loader = DataLoader(single_input_dataset,39                                            batch_size=1,40                                            shuffle=False,41                                            num_workers=2)42 43        # Perform inference on the single input44        output = inference_fn(single_input_loader, model, device)45 46        prediction = output.reshape((1, CFG.max_len))47        char_probs = get_char_probs([context], prediction, tokenizer)48        predictions = np.mean([char_probs], axis=0)49        results = get_results(predictions, th=0.5)50 51        print(results)52        return get_text(context, results[0])53    54    inputs = [gr.inputs.Textbox(label="Context Para", lines=10), gr.inputs.Textbox(label="Question", lines=1)]55    output = gr.outputs.Textbox(label="Answer")56 57    app = gr.Interface(fn=get_answer, inputs=inputs, outputs=output, allow_flagging='never')58 59    app.launch()60    print(get_answer(inputs[0], inputs[1]))61