MienOlle commited on
Commit
fcd3838
·
1 Parent(s): 1f07ce2

Upload main py

Browse files
Files changed (1) hide show
  1. main.py +23 -0
main.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from huggingface_hub import hf_hub_download
2
+ import torch
3
+ from transformers import AutoModelForSequenceClassification as modelSC, AutoTokenizer as token
4
+
5
+ model_path = hf_hub_download(repo_id="MienOlle/sentiment_analysis_api",
6
+ filename="sentimentAnalysis.pth"
7
+ )
8
+ modelToken = token.from_pretrained("mdhugol/indonesia-bert-sentiment-classification")
9
+ model = modelSC.from_pretrained("mdhugol/indonesia-bert-sentiment-classification", num_labels=3)
10
+ model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
11
+ model.eval()
12
+
13
+ def predict(input):
14
+ inputs = modelToken(input, return_tensors="pt", padding=True, truncation=True, max_length=512)
15
+
16
+ with torch.no_grad():
17
+ outputs = model(**inputs)
18
+
19
+ logits = outputs.logits
20
+ ret = logits.argmax.item()
21
+
22
+ labels = ["positive", "neutral", "negative"]
23
+ return labels[ret]