Views
No views yet
1from transformers import pipeline
2
3home_path = os.getcwd()
4print(home_path)
5# loading model
6model_path = f"{home_path}/checkpoint-326040"
7assert os.path.isdir(model_path)
8model = AutoModelForSequenceClassification.from_pretrained(model_path)
9tokenizer = AutoTokenizer.from_pretrained(model_path)
10# loading your data, formatted like ['product_name']
11val_df = pd.read_csv(f"your_data.csv")
12# loading dictionary convert label to category
13with open(f"{home_path}/label_name_6country.json", "r") as f:
14 cat_dic = json.load(f)
15label_name_dict = cat_dic["label_name_dict"]
16name_label_dict = cat_dic["name_label_dict"]
17# loading classifier pipeline
18classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, device="cuda:0")
19def get_label(name):
20 with torch.no_grad():
21 labels = classifier(product_names[1], top_k=2)
22 label = labels[0]["label"].split('_')[1]
23 if str(label_name_dict[label]) == "0":
24 label = labels[1]["label"].split('_')[1]
25 return label_name_dict[label]
26# predict the label
27def predict_df(df):
28 df["predict_label"] = df["product_name"].apply(lambda x: get_label(x))
29 return df