clean up audio caches automatically (#18)

- clean up audio caches automatically (d69879ca9279d55ddad0b7a06e499cd6e4f2e81a)

app.py

@@ -5,6 +5,7 @@ from datetime import datetime
 import time
 import json
 import logging
+from logging.handlers import RotatingFileHandler
 from typing import List, Dict, Tuple
 import gradio as gr
 
@@ -15,6 +16,29 @@ import soundfile as sf
 from huggingface_hub import CommitScheduler
 from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer
 
+class MyCommitScheduler(CommitScheduler):
+    def __init__(self, *args, cleanup_dir: str, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.cleanup_dir = cleanup_dir
+
+    def push_to_hub(self):
+        super().push_to_hub()
+        self.clean_up()
+        logger.info(f"Clean up directory: {self.cleanup_dir}.")
+
+    def clean_up(self):
+        if os.path.isdir(self.cleanup_dir):
+            for filename in os.listdir(self.cleanup_dir):
+                file_path = os.path.join(self.cleanup_dir, filename)
+                try:
+                    if os.path.isfile(file_path) or os.path.islink(file_path):
+                        os.remove(file_path)
+                    elif os.path.isdir(file_path):
+                        shutil.rmtree(file_path)
+                except Exception as e:
+                    logger.error(f"{e}")
+
+
 def setup_logger(log_dir: str, log_file: str, level: int = logging.INFO):
     os.makedirs(log_dir, exist_ok=True)
 
@@ -32,8 +56,11 @@ def setup_logger(log_dir: str, log_file: str, level: int = logging.INFO):
     console_handler.setFormatter(formatter)
     logger.addHandler(console_handler)
 
-
-    file_handler = logging.FileHandler(os.path.join(log_dir, log_file), encoding="utf-8")
+    file_handler = RotatingFileHandler(
+        os.path.join(log_dir, log_file),
+        encoding="utf-8",
+        maxBytes= 5 * 1024 * 1024, # max size is 5 MB
+        backupCount=1)
     file_handler.setFormatter(formatter)
     logger.addHandler(file_handler)
 
@@ -64,7 +91,7 @@ def infer(text_input: str, audio_input: str, sr: int = 16000) -> str:
         logger.error("Invalid inputs! Please check your prompt or reupload the audio file.")
         raise gr.Error("Invalid inputs! Please check your prompt or reupload the audio file.")
     
-    audio_path, request_id = handle_request(text_input, audio_input, sr=sr)
+    audio_path, request_id = handle_request(text_input, audio_input, audio_dir=cleanup_dir, sr=sr)
     message = construct_message(text_input, audio_path)
     now = datetime.now(time_zone)
     now_timestamp = now.timestamp()
@@ -74,7 +101,10 @@ def infer(text_input: str, audio_input: str, sr: int = 16000) -> str:
     # response = _infer_test(message)
     end_time = time.perf_counter()
     cost_time = round(end_time - start_time, 2)
-    log_data_str = handle_logging(text_input, audio_path, response, request_id, now_timestamp, cost_time)
+    log_data_str = handle_logging(
+        text_input, audio_path, response, 
+        request_id, now_str, cost_time,
+        log_dir=log_dir, log_file=response_log_file)
     logger.info(log_data_str)
     gr.Info(
         message=f"Inference's been done, took {cost_time}s.",
@@ -106,24 +136,23 @@ def is_valid_inputs(text_input: str, audio_input: str) -> bool:
 
 def handle_request(
         text_input: str, audio_input: str, 
-        usr_dir: str = "resources/usr", sr: int = 16000):
+        audio_dir: str = "resources/logs/audio", sr: int = 16000):
     request_id = os.urandom(16).hex()
-    out_dir = os.path.join(usr_dir, request_id)
-    os.makedirs(out_dir, exist_ok=True)
+    os.makedirs(audio_dir, exist_ok=True)
 
-    audio_path = os.path.join(out_dir, f"{request_id}.wav")
+    audio_path = os.path.join(audio_dir, f"{request_id}.wav")
     try:
         audio = preprocess_audio(audio_input, target_sr=sr)
         sf.write(audio_path, audio, sr, format="WAV")
 
         return audio_path, request_id
     except Exception as e:
-        shutil.rmtree(out_dir, ignore_errors=True)
+        shutil.rmtree(audio_dir, ignore_errors=True)
         raise gr.Error(f"{e}")
 
 def handle_logging(
         text_input: str, audio_path: str, response: str,
-        request_id: str, timestamp: float, cost_time: float,
+        request_id: str, timestamp: str, cost_time: float,
         log_dir: str = "resources/logs", log_file : str = "response.log"):
     log_file_path = os.path.join(log_dir, log_file)
     audio_log_dir = os.path.join(log_dir, "audio")
@@ -140,11 +169,9 @@ def handle_logging(
     log_data_str = json.dumps(log_data, ensure_ascii=False)
     if log_scheduler is not None:
         with log_scheduler.lock:
-            shutil.copy2(audio_path, audio_log_path)
             with open(log_file_path, mode='a', encoding="utf-8") as writer:
                 writer.write(log_data_str + "\n")
     else:
-        shutil.copy2(audio_path, audio_log_path)
         with open(log_file_path, mode='a', encoding="utf-8") as writer:
             writer.write(log_data_str + "\n")       
 
@@ -182,11 +209,13 @@ def construct_message(
 def create_log_scheduler(
         repo_id: str, repo_type: str, 
         folder_path: str, path_in_repo: str,
-        every: int = 3600):
-    scheduler = CommitScheduler(
+        every: int = 30,
+        cleanup_dir: str = "resources/logs/audio"):
+    scheduler = MyCommitScheduler(
         repo_id=repo_id, repo_type=repo_type,
         folder_path=folder_path, path_in_repo=path_in_repo,
-        every=every)
+        every=every,
+        cleanup_dir=cleanup_dir)
     return scheduler
 
 def update_audio_input(choice):
@@ -208,9 +237,15 @@ def enable_button():
 
 if __name__ == "__main__":
     time_zone = ZoneInfo("Asia/Shanghai")
+    appstart_time = datetime.now(time_zone)
+    appstart_timestamp = appstart_time.timestamp()
+    appstart_str = appstart_time.strftime("%Y-%m-%d_%H-%M-%S")
+
     model_name = "mispeech/MiDashengLM-7B"
     json_path = "resources/examples.json"
     log_dir = "resources/logs"
+    cleanup_dir = "resources/logs/audio" 
+    response_log_file = f"response_{appstart_str}.log"
     log_file = "gradio.log" # DEBUG logs
 
     model, processor, tokenizer = load_model(model_name)
@@ -220,7 +255,8 @@ if __name__ == "__main__":
     log_scheduler = create_log_scheduler(
         repo_id="mispeech/MiDashengLM-logs", repo_type="dataset",
         folder_path="resources/logs", path_in_repo="logs", 
-        every=15) # commit every 15 minute
+        every=15,
+        cleanup_dir=cleanup_dir) # commit every 15 minute
 
     with gr.Blocks() as demo:
         gr.Markdown("#🪄 Select an example or upload your own audio")