Update app.py

app.py

@@ -53,6 +53,14 @@ Categorize the website into one of the 3 categories:
         outputs = model.generate(**inputs, max_new_tokens=64, use_cache=True)
         ans = tokenizer.batch_decode(outputs)[0]
         ans_pred = ans.split('### Response:')[1].split('<')[0]
+        
+        if 'OTHER' in ans_pred:
+            ans_pred = 'OTHER'
+        elif 'NEWS/BLOG' in ans_pred:
+            ans_pred = 'NEWS/BLOG'
+        elif 'E-commerce' in ans_pred:
+            ans_pred = 'E-commerce'
+            
         return ans_pred
 
     except Exception as e: