Update README.md

README.md

@@ -30,6 +30,41 @@ This is the model card of a 🤗 transformers model that has been pushed on the
 
 ## How to Get Started with the Model
 
+class IntentClassifier:
+    def __init__(self, model_name="serj/intent-classifier", device="cuda"):
+        self.model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
+        self.tokenizer = T5Tokenizer.from_pretrained(model_name)
+        self.device = device
+
+
+def build_prompt(text, prompt="", company_name="", company_specific=""):
+    if company_name == "Pizza Mia":
+        company_specific = "This company is a pizzeria place."
+    if company_name == "Online Banking":
+        company_specific = "This company is an online banking."
+
+    return f"Company name: {company_name} is doing: {company_specific}\nCustomer: {text}.\nEND MESSAGE\nChoose one topic that matches customer's issue.\n{prompt}\nClass name: "
+
+
+def predict(self, text, prompt_options, company_name, company_portion) -> str:
+    input_text = build_prompt(text, prompt_options, company_name, company_portion)
+    # print(input_text)
+    # Tokenize the concatenated inp_ut text
+    input_ids = self.tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True).to(self.device)
+
+    # Generate the output
+    output = self.model.generate(input_ids)
+
+    # Decode the output tokens
+    decoded_output = self.tokenizer.decode(output[0], skip_special_tokens=True)
+
+    return decoded_output
+
+
+m = IntentClassifier("serj/intent-classifier")
+print(m.predict("Hey, after recent changes, I want to cancel subscription, please help.",
+                "OPTIONS:\n refund\n cancel subscription\n damaged item\n return item\n", "Company",
+                "Products and subscriptions"))
 
 
 [More Information Needed]