Views
No views yet
1from transformers import T5ForConditionalGeneration, T5Tokenizer
2
3model = T5ForConditionalGeneration.from_pretrained("{repo_id}")
4tokenizer = T5Tokenizer.from_pretrained("{repo_id}")
5
6def predict(text):
7 prompt = f"Predict the product category and subcategory in the following format: 'Category: <CATEGORY> | Subcategory: <SUBCATEGORY>'. Product: {text}"
8 inputs = tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True)
9
10 outputs = model.generate(**inputs, max_length=32, num_beams=4)
11 return tokenizer.decode(outputs[0], skip_special_tokens=True)
12
13# Example
14result = predict("Pantene Suave & Liso Shampoo")
15print(result)