model_2 / train.py
Junnyfirst34's picture
Create train.py
2f4b776 verified
from transformers import AutoModelForSequenceClassification, AutoTokenizer
def main():
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(
"bert-base-uncased", num_labels=3
)
text = "pick up the object"
inputs = tokenizer(text, return_tensors="pt")
outputs = model(**inputs)
print(outputs.logits)
if __name__ == "__main__":
main()