ams-core-api / core /classifier.py
Amssou's picture
Create classifier.py
94463f6 verified
Raw
History Blame Contribute Delete
1.23 kB
# core/classifier.py
class PromptClassifier:
def __init__(self):
self.business_keywords = [
"startup", "business", "marché", "client",
"produit", "vente", "rentable", "entreprise"
]
self.tech_keywords = [
"ia", "intelligence artificielle",
"algorithme", "code", "développer"
]
self.social_keywords = [
"tiktok", "instagram", "réseaux",
"contenu", "publication", "vidéo"
]
self.web_keywords = [
"actualité", "actuel", "aujourd",
"président", "news", "prix",
"bitcoin", "bourse", "météo"
]
self.zero_budget_keywords = [
"0 budget", "zéro", "sans argent",
"gratuit"
]
def classify(self, prompt: str) -> dict:
p = prompt.lower()
return {
"business": any(k in p for k in self.business_keywords),
"tech": any(k in p for k in self.tech_keywords),
"social": any(k in p for k in self.social_keywords),
"actualité": any(k in p for k in self.web_keywords),
"zero_budget": any(k in p for k in self.zero_budget_keywords),
}