davanstrien HF staff commited on
Commit
572201a
·
verified ·
1 Parent(s): 7243626

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -5
app.py CHANGED
@@ -6,7 +6,9 @@ import logging
6
  from typing import Tuple, Literal
7
  import functools
8
  import spaces
9
- from cashews import cache
 
 
10
  # Set up logging
11
  logging.basicConfig(level=logging.INFO)
12
  logger = logging.getLogger(__name__)
@@ -16,8 +18,8 @@ MODEL_NAME = "davanstrien/Smol-Hub-tldr"
16
  model = None
17
  tokenizer = None
18
  device = None
19
-
20
- cache.setup("mem://", size_limit="4gb")
21
 
22
  def load_model():
23
  global model, tokenizer, device
@@ -87,9 +89,9 @@ def _generate_summary_gpu(card_text: str, card_type: str) -> str:
87
 
88
  return summary
89
 
90
- @cache(ttl="6h")
91
  def generate_summary(card_text: str, card_type: str) -> str:
92
- """Cached wrapper for generate_summary."""
93
  return _generate_summary_gpu(card_text, card_type)
94
 
95
  def summarize(hub_id: str = "", card_type: str = "model", content: str = "") -> str:
 
6
  from typing import Tuple, Literal
7
  import functools
8
  import spaces
9
+ from cachetools import TTLCache
10
+ from cachetools.func import ttl_cache
11
+ import time
12
  # Set up logging
13
  logging.basicConfig(level=logging.INFO)
14
  logger = logging.getLogger(__name__)
 
18
  model = None
19
  tokenizer = None
20
  device = None
21
+ CACHE_TTL = 6 * 60 * 60 # 6 hours in seconds
22
+ CACHE_MAXSIZE = 100
23
 
24
  def load_model():
25
  global model, tokenizer, device
 
89
 
90
  return summary
91
 
92
+ @ttl_cache(maxsize=CACHE_MAXSIZE, ttl=CACHE_TTL)
93
  def generate_summary(card_text: str, card_type: str) -> str:
94
+ """Cached wrapper for generate_summary with TTL."""
95
  return _generate_summary_gpu(card_text, card_type)
96
 
97
  def summarize(hub_id: str = "", card_type: str = "model", content: str = "") -> str: