henhenhahi111112/en_th

How to use

home_path = os.getcwd()
print(home_path)

# loading model
model_path = f"{home_path}/cate_calssification_en_thai_checkpoint_150320"
assert os.path.isdir(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)


# loading your data, formatted like ['product_name']
val_df = pd.read_csv(f"your_data.csv")

# loading dictionary convert label to category 
with open(f"{home_path}/label_name.json", "r") as f:
    cat_dic = json.load(f)
label_name_dict = cat_dic["label_name_dict"]
name_label_dict = cat_dic["name_label_dict"]

# loading classifier pipeline
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, device="cuda:0")
def get_label(name):
    label = classifier(name)[0]["label"].split('_')[1]
    return label_name_dict[label]

# predict the label
def predict_df(df):
    df["predict_label"] = df["product_name"].apply(lambda x: get_label(x))
    return df