52 lines
2.3 KiB
Python
52 lines
2.3 KiB
Python
import json
|
|
from app.config import settings
|
|
|
|
CATEGORIES = ["Politics","National","International","Sports","Business","Technology",
|
|
"Education","Health","Science","Entertainment","Crime","Court","Weather",
|
|
"Editorial","Opinion","Local News","State News","Environment","Lifestyle",
|
|
"Automobile","Travel","Agriculture","Jobs","Classifieds","Others"]
|
|
|
|
SUBCATS = {"Sports": ["Cricket","Football","Tennis","Others"],
|
|
"Business": ["Finance","Markets","Economy"],
|
|
"Technology": ["AI","Gadgets","Software"],
|
|
"Entertainment": ["Movies","TV","Music"],
|
|
"Politics": ["State Politics","National Politics","International"]}
|
|
|
|
PROMPT = """You are a news classifier. Given the headline and text, return strict JSON:
|
|
{{"category": one of {cats}, "subcategory": string, "confidence": 0-1,
|
|
"summary": 2-sentence summary, "keywords": [up to 8],
|
|
"persons": [], "locations": [], "organizations": [], "dates": [], "sentiment": "pos|neu|neg"}}
|
|
|
|
HEADLINE: {headline}
|
|
TEXT: {text}
|
|
"""
|
|
|
|
class Categorizer:
|
|
def __init__(self):
|
|
self.provider = settings.LLM_PROVIDER
|
|
if self.provider == "gemini":
|
|
import google.generativeai as genai
|
|
genai.configure(api_key=settings.LLM_API_KEY)
|
|
self.model = genai.GenerativeModel("gemini-1.5-flash")
|
|
|
|
def classify(self, headline, text):
|
|
prompt = PROMPT.format(cats=CATEGORIES, headline=headline, text=text[:4000])
|
|
try:
|
|
if self.provider == "gemini":
|
|
resp = self.model.generate_content(prompt)
|
|
raw = resp.text.strip().strip("```json").strip("```")
|
|
return json.loads(raw)
|
|
except Exception as e:
|
|
return self._fallback(headline, text)
|
|
|
|
def _fallback(self, headline, text):
|
|
# Zero-shot via local model if LLM unavailable
|
|
from transformers import pipeline
|
|
clf = pipeline("zero-shot-classification",
|
|
model="facebook/bart-large-mnli")
|
|
res = clf(headline + ". " + text[:500], CATEGORIES)
|
|
return {"category": res["labels"][0], "subcategory": "",
|
|
"confidence": res["scores"][0], "summary": "", "keywords": [],
|
|
"persons": [], "locations": [], "organizations": [],
|
|
"dates": [], "sentiment": "neu"}
|