Junnyfirst34 commited on
Commit
2f4b776
·
verified ·
1 Parent(s): 411a1ea

Create train.py

Browse files
Files changed (1) hide show
  1. train.py +16 -0
train.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer
2
+
3
+ def main():
4
+ tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
5
+ model = AutoModelForSequenceClassification.from_pretrained(
6
+ "bert-base-uncased", num_labels=3
7
+ )
8
+
9
+ text = "pick up the object"
10
+ inputs = tokenizer(text, return_tensors="pt")
11
+
12
+ outputs = model(**inputs)
13
+ print(outputs.logits)
14
+
15
+ if __name__ == "__main__":
16
+ main()