Are you the author?
GPU deploy isn't supported for this model type.
1import json
2import sys
3import logging
4import torch
5from torch import nn
6from transformers import ElectraConfig
7from transformers import ElectraModel, AutoTokenizer, ElectraTokenizer, ElectraForSequenceClassification
8
9logging.basicConfig(
10 level=logging.INFO,
11 format='[{%(filename)s:%(lineno)d} %(levelname)s - %(message)s',
12 handlers=[
13 logging.FileHandler(filename='tmp.log'),
14 logging.StreamHandler(sys.stdout)
15 ]
16)
17logger = logging.getLogger(__name__)
18
19max_seq_length = 128
20classes = ['Neg', 'Pos']
21
22tokenizer = AutoTokenizer.from_pretrained("daekeun-ml/koelectra-small-v3-nsmc")
23device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
24
25
26def model_fn(model_path=None):
27 ####
28 # If you have your own trained model
29 # Huggingface pre-trained model: 'monologg/koelectra-small-v3-discriminator'
30 ####
31 #config = ElectraConfig.from_json_file(f'{model_path}/config.json')
32 #model = ElectraForSequenceClassification.from_pretrained(f'{model_path}/model.pth', config=config)
33
34 # Download model from the Huggingface hub
35 model = ElectraForSequenceClassification.from_pretrained('daekeun-ml/koelectra-small-v3-nsmc')
36 model.to(device)
37 return model
38
39
40def input_fn(input_data, content_type="application/jsonlines"):
41 data_str = input_data.decode("utf-8")
42 jsonlines = data_str.split("\n")
43 transformed_inputs = []
44
45 for jsonline in jsonlines:
46 text = json.loads(jsonline)["text"][0]
47 logger.info("input text: {}".format(text))
48 encode_plus_token = tokenizer.encode_plus(
49 text,
50 max_length=max_seq_length,
51 add_special_tokens=True,
52 return_token_type_ids=False,
53 padding="max_length",
54 return_attention_mask=True,
55 return_tensors="pt",
56 truncation=True,
57 )
58 transformed_inputs.append(encode_plus_token)
59
60 return transformed_inputs
61
62
63def predict_fn(transformed_inputs, model):
64 predicted_classes = []
65
66 for data in transformed_inputs:
67 data = data.to(device)
68 output = model(**data)
69
70 softmax_fn = nn.Softmax(dim=1)
71 softmax_output = softmax_fn(output[0])
72 _, prediction = torch.max(softmax_output, dim=1)
73
74 predicted_class_idx = prediction.item()
75 predicted_class = classes[predicted_class_idx]
76 score = softmax_output[0][predicted_class_idx]
77 logger.info("predicted_class: {}".format(predicted_class))
78
79 prediction_dict = {}
80 prediction_dict["predicted_label"] = predicted_class
81 prediction_dict['score'] = score.cpu().detach().numpy().tolist()
82
83 jsonline = json.dumps(prediction_dict)
84 logger.info("jsonline: {}".format(jsonline))
85 predicted_classes.append(jsonline)
86
87 predicted_classes_jsonlines = "\n".join(predicted_classes)
88 return predicted_classes_jsonlines
89
90
91def output_fn(outputs, accept="application/jsonlines"):
92 return outputs, accept1>>> from inference_nsmc import model_fn, input_fn, predict_fn, output_fn
2>>> with open('samples/nsmc.txt', mode='rb') as file:
3>>> model_input_data = file.read()
4>>> model = model_fn()
5>>> transformed_inputs = input_fn(model_input_data)
6>>> predicted_classes_jsonlines = predict_fn(transformed_inputs, model)
7>>> model_outputs = output_fn(predicted_classes_jsonlines)
8>>> print(model_outputs[0])
9
10[{inference_nsmc.py:47} INFO - input text: 이 영화는 최고의 영화입니다
11[{inference_nsmc.py:47} INFO - input text: 최악이에요. 배우의 연기력도 좋지 않고 내용도 너무 허접합니다
12[{inference_nsmc.py:77} INFO - predicted_class: Pos
13[{inference_nsmc.py:84} INFO - jsonline: {"predicted_label": "Pos", "score": 0.9619030952453613}
14[{inference_nsmc.py:77} INFO - predicted_class: Neg
15[{inference_nsmc.py:84} INFO - jsonline: {"predicted_label": "Neg", "score": 0.9994170665740967}
16{"predicted_label": "Pos", "score": 0.9619030952453613}
17{"predicted_label": "Neg", "score": 0.9994170665740967}{"text": ["이 영화는 최고의 영화입니다"]}
{"text": ["최악이에요. 배우의 연기력도 좋지 않고 내용도 너무 허접합니다"]}