Spaces:
Sleeping
Sleeping
File size: 2,872 Bytes
9db14d3 |
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 |
import os
import sys
sys.path.insert(0, os.path.abspath('./'))
import torch
from tqdm.auto import tqdm
from torch.utils.data import DataLoader, random_split
from transformers import AutoTokenizer, AutoModelForMaskedLM
from event_detection_dataset import *
from event_detection_model import *
import gradio as gr
#print(f"Gradio version: {gr.__version__}")
def predict(data):
data=[data]
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#print(f"Device {device}")
"""Load Tokenizer"""
tokenizer = AutoTokenizer.from_pretrained('distilbert-base-cased', use_fast=True)
"""Tokenized Inputs"""
tokenized_inputs = tokenizer(
data,
add_special_tokens=True,
max_length=512,
padding='max_length',
return_token_type_ids=True,
truncation=True,
is_split_into_words=True
)
"""Load Model"""
model_path = "./"
#print("model_path:", model_path)
#print("================ load model ===========================")
model = DistillBERTClass('distilbert-base-cased')
#print("================ model init ===========================")
pretrained_model=torch.load(model_path + "/event_domain_final.pt",map_location=torch.device('cpu'))
model.load_state_dict(pretrained_model['model_state_dict'])
model.to(device)
"""Make Prediction"""
model.eval()
ids = torch.tensor(tokenized_inputs['input_ids']).to(device)
mask = torch.tensor(tokenized_inputs['attention_mask']).to(device)
with torch.no_grad():
outputs = model(ids, mask)
max_val, max_idx = torch.max(outputs.data, dim=1)
#print("=============== inference result =================")
#print(f"predicted class {max_idx}")
id2tags={0: "Acquisition",1: "I-Positive Clinical Trial & FDA Approval",2: "Dividend Cut",3: "Dividend Increase",4: "Guidance Increase",5: "New Contract",6: "Dividend",7: "Reverse Stock Split",8: "Special Dividend ",9: "Stock Repurchase",10: "Stock Split",11: "Others"}
return id2tags[max_idx.item()]
title="Financial Event Detection"
description="Predict Finacial Events."
article="modified the model in the following paper: Zhou, Z., Ma, L., & Liu, H. (2021)."
example_list=[["Investors who receive dividends can choose to take them as cash or as additional shares."]]
# Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
inputs="text", # what are the inputs?
outputs="text", # our fn has two outputs, therefore we have two outputs
examples=example_list,
title=title,
description=description,
article=article)
# Launch the demo!
demo.launch(debug=False, share=True)
|