thak123 commited on
Commit
9251dbc
1 Parent(s): 2de6916

Create model.py

Browse files
Files changed (1) hide show
  1. model.py +19 -0
model.py ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import transformers
2
+ from mtm import MultitaskModel
3
+ from config import model_name, drop_out
4
+
5
+ multitask_model = MultitaskModel.create(
6
+ model_name=model_name,
7
+ model_type_dict={
8
+ "document": transformers.AutoModelForSequenceClassification,
9
+ "paragraph": transformers.AutoModelForSequenceClassification,
10
+ "sentence": transformers.AutoModelForSequenceClassification,
11
+ },
12
+ model_config_dict={
13
+ "document": transformers.AutoConfig.from_pretrained(model_name, num_labels=3, hidden_dropout_prob=drop_out, attention_probs_dropout_prob=drop_out),
14
+ "paragraph": transformers.AutoConfig.from_pretrained(model_name, num_labels=3, hidden_dropout_prob=drop_out, attention_probs_dropout_prob=drop_out),
15
+ "sentence": transformers.AutoConfig.from_pretrained(model_name, num_labels=3, hidden_dropout_prob=drop_out, attention_probs_dropout_prob=drop_out),
16
+ },
17
+ )
18
+
19
+ tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)