services/article-parser/src/controllers/clustering.py - angular - boris00043/semantic-papers

services/article-parser/src/controllers/clustering.py
112 строк4 KB
Бредихин Борис Андреевич
repredict problems and tasks
23 авг 2025, 14:36
23 авг 2025, 14:364ab4b27
from collections import defaultdict
import datetime
import logging
from os import PathLike
import os

from arango import CursorEmptyError

from ai.text_clusterizer import TextClusterizer
from common import ollama_client
from db.arango import ArangoDatabase
from schemas import ClusteringModel


class ProblemsTasksClusterizingController:
    def __init__(self, articles_db: ArangoDatabase, clustering_models_dir: PathLike):
        self.articles_db: ArangoDatabase = articles_db
        self.clustering_models_dir: PathLike = clustering_models_dir
        self.logger = logging.getLogger("ProblemsTasksClusterizingController")
        self.clustering_models_collection = self.articles_db.get_vertex_collection(
            "ClusteringModel"
        )

    def train_clustering_model(self, model_name: str, model_object: TextClusterizer):
        self.logger.info(
            "Training clustering model", extra={"clustering_model_name": model_name}
        )

        dataset = self.get_problems_and_tasks()
        model_object.fit(dataset)

        self.logger.info(
            "Saving clustering model", extra={"clustering_model_name": model_name}
        )

        model_object.save(os.path.join(self.clustering_models_dir, model_name))
        self.clustering_models_collection.insert(
            ClusteringModel(
                model_name=model_name, created_at=datetime.datetime.now().isoformat()
            )
        )

    def get_problems_and_tasks(self) -> list[str]:
        return [
            i["description"]
            for i in self.articles_db.get_vertex_collection("ProblemsAndTasks").all()
        ]

    def list_clustering_models(self) -> list[ClusteringModel]:
        return list(self.clustering_models_collection.all())

    def load_clustering_model(self, model_name: str) -> TextClusterizer:
        if model_name not in os.listdir(self.clustering_models_dir):
            raise ValueError("Model not found")
        else:
            return TextClusterizer.load(
                os.path.join(self.clustering_models_dir, model_name), ollama_client
            )

    def predict_theme_cluster(self, model_name: str, text: str) -> str:
        if model_name not in os.listdir(self.clustering_models_dir):
            raise ValueError("Model not found")
        else:
            model = self.load_clustering_model(model_name)
            return model.predict(text)

    def predict_theme_cluster_db(self, model_name: str, problem_task_id: str) -> str:
        db_obj = self.articles_db.get_vertex_collection("ProblemsAndTasks").get(
            problem_task_id
        )
        return self.predict_theme_cluster(model_name, db_obj["description"])

    def update_all(self, model_name: str):
        for problem_task in self.articles_db.get_vertex_collection(
            "ProblemsAndTasks"
        ).all():
            problem_task["theme_cluster"] = self.predict_theme_cluster_db(
                model_name, problem_task["_key"]
            )
            problem_task["cluster_model"] = model_name
            self.articles_db.get_vertex_collection("ProblemsAndTasks").update(
                problem_task
            )

    def get_all(self):
        result = defaultdict(list)
        for a in self.articles_db.get_vertex_collection("ProblemsAndTasks").all():
            theme_cluster = a.get("theme_cluster", "<<UNPROCESSED>>")
            try:
                article_id = (
                    self.articles_db.get_edge_collection(
                        "article_to_problems_and_tasks",
                        ["Article"],
                        ["ProblemsAndTasks"],
                    )
                    .find(
                        {
                            "_to": a["_id"],
                        }
                    )
                    .pop()["_from"]
                )
            except CursorEmptyError:
                self.logger.warning(
                    "Article not found in article_to_problems_and_tasks",
                    extra={"article_id": a["_id"]},
                )
                continue
            result[theme_cluster].append(
                {"description": a["description"], "article_id": article_id}
            )
        return result