Team Ai
Apppublic

hishn12/Email_Classification_without_LLM

sourceHugging Facemitupdated 1y agoView on Hugging Face
1likes
cache_manager.py80 linesDownload Raw Back to root
1import os2from typing import Dict, Any, Optional3import json, hashlib4from datetime import datetime, timedelta5from logger import setup_logger6 7logger = setup_logger(__name__)8 9 10class CacheManager:11    """Manages caching of masked text and PII mappings."""12 13    def __init__(self, cache_file: str = "cache/pii_cache.json"):14        self.cache_file = cache_file15        os.makedirs(os.path.dirname(self.cache_file), exist_ok=True)  # Ensure directory exists16        self.cache: Dict[str, Any] = {}17        self.load_cache()18 19    def add_to_cache(20            self,21            original: str,22            masked: str,23            findings: Dict,24            ttl_hours: int = 2425    ):26        """Add entry to cache with TTL."""27        try:28            self.cache[self._hash(original)] = {29                'masked': masked,30                'findings': findings,31                'expires': (32                        datetime.now() + timedelta(hours=ttl_hours)33                ).isoformat()34            }35            self.save_cache()36 37        except Exception as e:38            logger.error(f"Cache addition error: {e}")39 40    def get_from_cache(self, text: str) -> Optional[Dict]:41        """Retrieve cached entry if exists and not expired."""42        try:43            key = self._hash(text)44            if key in self.cache:45                entry = self.cache[key]46                if datetime.now() < datetime.fromisoformat(entry['expires']):47                    return entry48                else:49                    del self.cache[key]50                    self.save_cache()51            return None52 53        except Exception as e:54            logger.error(f"Cache retrieval error: {e}")55            return None56 57    def save_cache(self):58        """Save cache to file."""59        try:60            with open(self.cache_file, 'w') as f:61                json.dump(self.cache, f, indent=2)62        except Exception as e:63            logger.error(f"Cache save error: {e}")64 65    def load_cache(self):66        """Load cache from file."""67        try:68            with open(self.cache_file, 'r') as f:69                self.cache = json.load(f)70        except FileNotFoundError:71            self.cache = {}72        except Exception as e:73            logger.error(f"Cache load error: {e}")74            self.cache = {}75 76    @staticmethod77    def _hash(text: str) -> str:78        """Create hash for cache key."""79        return hashlib.md5(text.encode()).hexdigest()80