tweaks to beats config/init (#114)

<!-- av pr metadata
This information is embedded by the av CLI when creating PRs to track
the status of stacks when using Aviator. Please do not delete or edit
this section of the PR.
```
{"parent":"main","parentHead":"","trunk":"main"}
```
-->

Dockerfile

@@ -11,6 +11,7 @@ COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
 RUN apt-get update && apt-get install -y \
     git \
     git-lfs \
+    ffmpeg \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/* \
     && git lfs install
@@ -18,7 +19,7 @@ RUN apt-get update && apt-get install -y \
 # TODO: Pin esp-research and esp-data revisions
 # TODO: remove hf-app branch once merged
 RUN --mount=type=secret,id=GH_TOKEN,mode=0444,required=true \
-    git clone -b hf-app --single-branch --depth 1 https://$(cat /run/secrets/GH_TOKEN)@github.com/earthspecies/esp-research.git /app/esp-research && \
+    git clone -b update-hf-app --single-branch --depth 1 https://$(cat /run/secrets/GH_TOKEN)@github.com/earthspecies/esp-research.git /app/esp-research && \
     git clone --single-branch --depth 1 https://$(cat /run/secrets/GH_TOKEN)@github.com/earthspecies/esp-data.git /app/esp-data
 
 # esp-research installs esp-data from gcloud artifact registry, which is not

app.py

@@ -3,6 +3,7 @@ import uuid
 from pathlib import Path
 
 import gradio as gr
+import librosa
 import matplotlib.pyplot as plt
 import numpy as np
 import soundfile as sf
@@ -11,49 +12,34 @@ import torch
 import torchaudio
 
 from esp_research.logging import logger
-from hub_logger import upload_data
-
-# from NatureLM.infer import Pipeline
-# from NatureLM.models.NatureLM import NatureLM
-from naturelm_audio import NatureLM  # noqa: F401
+from hub_logger import log_interaction
+from naturelm_audio import GenerationConfig, NatureLM
 
 APP_DIR = Path(__file__).resolve().parent
 STATIC_DIR = APP_DIR / "static"
 ASSETS_DIR = APP_DIR / "assets"
 
-SAMPLE_RATE = 16000  # Default sample rate for NatureLM-audio
+# TODO: Set these values carefully later.
+SAMPLE_RATE = 16000
 MIN_AUDIO_DURATION: float = 0.5  # seconds
-MAX_HISTORY_TURNS = 3  # Maximum number of conversation turns to include in context (user + assistant pairs)
+MAX_AUDIO_DURATION: float = 60.0  # seconds
+MAX_HISTORY_TURNS = 3
 
-DEVICE: str = "cuda" if torch.cuda.is_available() else "cpu"
+assert torch.cuda.is_available(), "CUDA is required to run this app"
+DEVICE = "cuda"
 
 # TODO: derive model version from model metadata or config instead of hardcoding
-MODEL_VERSION = "1.5"
-
-
-class _MockModel:
-    """Placeholder model that returns dummy predictions."""
-
-    def __call__(
-        self,
-        audios: list[str],
-        queries: list[str],
-        **kwargs: object,
-    ) -> list[list[dict]]:
-        return [[{"prediction": "(mock) I don't know yet!"}] for _ in audios]
-
+MODEL_VERSION = "1.1"
+MODEL_REPO_ID = "EarthSpeciesProject/naturelm-audio-1.1.00-private"
 
-# TODO: replace with real model loading
+logger.info("Loading model from %s …", MODEL_REPO_ID)
+model = NatureLM.from_hf_hub(MODEL_REPO_ID)
+model = model.eval().to(DEVICE)
+logger.info("Model loaded successfully")
 
-# model = NatureLM.from_pretrained("EarthSpeciesProject/NatureLM-audio")
-# model = model.eval().to(DEVICE)
-# model = Pipeline(model)
-logger.info("Device: %s", DEVICE)
-model = _MockModel()
 
-
-def validate_audio_duration(audio_path: str) -> None:
-    """Validate that the audio file meets the minimum duration requirement.
+def validate_audio(audio_path: str) -> None:
+    """Validate that the audio file meets duration requirements.
 
     Parameters
     ----------
@@ -63,53 +49,18 @@ def validate_audio_duration(audio_path: str) -> None:
     Raises
     ------
     Error
-        If the audio duration is less than `MIN_AUDIO_DURATION`.
+        If the audio duration is outside [`MIN_AUDIO_DURATION`,
+        `MAX_AUDIO_DURATION`].
     """
     info = sf.info(audio_path)
-    duration = info.duration  # info.num_frames / info.sample_rate
+    duration = info.duration
     if duration < MIN_AUDIO_DURATION:
         raise gr.Error(f"Audio duration must be at least {MIN_AUDIO_DURATION} seconds.")
+    if duration > MAX_AUDIO_DURATION:
+        raise gr.Error(f"Audio duration must be at most {MAX_AUDIO_DURATION} seconds.")
 
 
 @spaces.GPU
-def prompt_lm(
-    audios: list[str],
-    queries: list[str] | str,
-    window_length_seconds: float = 10.0,
-    hop_length_seconds: float = 10.0,
-) -> list[str]:
-    """Generate response using the model.
-
-    Parameters
-    ----------
-    audios : list[str]
-        List of audio file paths.
-    queries : list[str] | str
-        Query or list of queries to process.
-    window_length_seconds : float
-        Length of the window for processing audio.
-    hop_length_seconds : float
-        Hop length for processing audio.
-
-    Returns
-    -------
-    list[list[dict]]
-        Nested list of prediction dictionaries for each audio-query pair.
-    """
-    if model is None:
-        return "❌ Model not loaded. Please check the model configuration."
-
-    with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
-        results: list[list[dict]] = model(
-            audios,
-            queries,
-            window_length_seconds=window_length_seconds,
-            hop_length_seconds=hop_length_seconds,
-            input_sample_rate=None,
-        )
-    return results
-
-
 def get_response(chatbot_history: list[dict], audio_input: str) -> list[dict]:
     """Generate response from the model based on user input and audio file.
 
@@ -134,56 +85,38 @@ def get_response(chatbot_history: list[dict], audio_input: str) -> list[dict]:
                 " Consider starting a new conversation with the Clear button."
             )
 
-        # Build conversation context from history
-        conversation_context = []
-        for message in chatbot_history:
-            if message["role"] == "user":
-                conversation_context.append(f"User: {message['content']}")
-            elif message["role"] == "assistant":
-                conversation_context.append(f"Assistant: {message['content']}")
-
-        # Get the last user message
-        last_user_message = ""
-        for message in reversed(chatbot_history):
-            if message["role"] == "user":
-                last_user_message = message["content"]
-                break
-
-        # Format the full prompt with conversation history
-        if len(conversation_context) > 2:  # More than just the current query
-            # Include previous turns (limit to last MAX_HISTORY_TURNS exchanges)
-            # recent_context = conversation_context[
-            #    -(MAX_HISTORY_TURNS + 1) : -1
-            # ]  # Exclude current message
-            recent_context = conversation_context
-
-            full_prompt = (
-                "Previous conversation:\n" + "\n".join(recent_context) + "\n\nCurrent question: " + last_user_message
+        # Load audio, mix to mono, and resample to model sample rate if needed
+        audio_np, sr = sf.read(audio_input, dtype="float32")
+        if audio_np.ndim > 1:
+            audio_np = np.mean(audio_np, axis=1)
+        if sr != SAMPLE_RATE:
+            audio_np = librosa.resample(
+                y=audio_np, orig_sr=sr, target_sr=SAMPLE_RATE, res_type="kaiser_best", scale=True
             )
-        else:
-            full_prompt = last_user_message
-
-        logger.debug("Full prompt with history: %s", full_prompt)
-
-        response = prompt_lm(
-            audios=[audio_input],
-            queries=[full_prompt.strip()],
-            window_length_seconds=100_000,
-            hop_length_seconds=100_000,
-        )
-        # get first item
-        if isinstance(response, list) and len(response) > 0:
-            response = response[0][0]["prediction"]
-            logger.info("Model response: %s", response)
-        else:
-            response = "No response generated."
+        audio_tensor = torch.from_numpy(audio_np).to(DEVICE)
+
+        # Build chat-format messages for model.generate().
+        # Gradio may return content as a list of parts on subsequent turns,
+        # so normalise to plain strings first.
+        messages: list[dict[str, str]] = []
+        for msg in chatbot_history:
+            text = msg["content"] if isinstance(msg["content"], str) else msg["content"][0]["text"]
+            if msg["role"] in ("user", "assistant"):
+                messages.append({"role": msg["role"], "content": text})
+
+        logger.debug("Messages: %s", messages)
+
+        response = model.generate(
+            audio=[audio_tensor],
+            messages=[messages],
+            generation_config=GenerationConfig(merging_alpha=0.7),
+        )[0]
+        logger.info("Model response: %s", response)
     except Exception as e:
         logger.exception("Error generating response: %s", e)
         response = "Error generating response. Please try again."
 
-    # Add model response to chat history
     chatbot_history.append({"role": "assistant", "content": response})
-
     return chatbot_history
 
 
@@ -269,7 +202,6 @@ def add_user_query(chatbot_history: list[dict], chat_input: str) -> list[dict]:
     list[dict]
         Updated chat history with the user message appended.
     """
-    # Validate input
     if not chat_input.strip():
         return chatbot_history
 
@@ -283,7 +215,7 @@ def log_to_hub(chatbot_history: list[dict], audio: str, session_id: str) -> None
         return
     user_text = chatbot_history[-2]["content"]
     model_response = chatbot_history[-1]["content"]
-    upload_data(audio, user_text, model_response, session_id, model_version=MODEL_VERSION)
+    log_interaction(audio, user_text, model_response, session_id, model_version=MODEL_VERSION)
 
 
 def main() -> gr.Blocks:
@@ -331,14 +263,6 @@ def main() -> gr.Blocks:
         with gr.Tabs():
             with gr.Tab("Analyze Audio"):
                 session_id = gr.State(str(uuid.uuid4()))
-                # uploaded_audio = gr.State()
-                # Status indicator
-                # status_text = gr.Textbox(
-                #     value=model_manager.get_status(),
-                #     label="Model Status",
-                #     interactive=False,
-                #     visible=True,
-                # )
 
                 with gr.Column(visible=True) as onboarding_message:
                     gr.HTML(
@@ -351,11 +275,11 @@ def main() -> gr.Blocks:
                         container=True,
                         interactive=True,
                         sources=["upload"],
+                        type="filepath",
                     )
-                    # check that audio duration is greater than MIN_AUDIO_DURATION
-                    # raise
+                    # Validate audio duration and sample rate on upload
                     audio_input.change(
-                        fn=validate_audio_duration,
+                        fn=validate_audio,
                         inputs=[audio_input],
                         outputs=[],
                     )
@@ -475,11 +399,6 @@ def main() -> gr.Blocks:
                     outputs=[plotter],
                 )
 
-                # When submit clicked first:
-                # 1. Validate and add user query to chat history
-                # 2. Get response from model
-                # 3. Clear the chat input box
-                # 4. Show clear button
                 chat_input.submit(
                     validate_and_submit,
                     inputs=[chatbot, chat_input],

hub_logger.py

@@ -1,3 +1,5 @@
+"""Log user interactions to a HuggingFace Hub dataset."""
+
 import json
 import os
 import uuid
@@ -8,59 +10,72 @@ from huggingface_hub import HfApi, HfFileSystem
 DATASET_REPO = "EarthSpeciesProject/naturelm-audio-space-logs"
 SPLIT = "test"
 TESTING = os.getenv("TESTING", "0") == "1"
-api = HfApi(token=os.getenv("HF_TOKEN", None))
-# Upload audio
-# check if file exists
-hf_fs = HfFileSystem(token=os.getenv("HF_TOKEN", None))
+
+_hf_token = os.getenv("HF_TOKEN", None)
+_api = HfApi(token=_hf_token)
+_fs = HfFileSystem(token=_hf_token)
+
+_METADATA_PATH = f"datasets/{DATASET_REPO}/{SPLIT}/metadata.jsonl"
 
 
-def upload_data(
+def log_interaction(
     audio: str | Path,
     user_text: str,
     model_response: str,
     session_id: str = "",
     model_version: str = "",
 ) -> None:
+    """Log a single user/model exchange (audio + text) to the Hub dataset.
+
+    Parameters
+    ----------
+    audio : str | Path
+        Local path to the audio file uploaded by the user.
+    user_text : str
+        The user's query text.
+    model_response : str
+        The model's generated response.
+    session_id : str
+        Unique identifier for the user session.
+    model_version : str
+        Version string of the model that produced the response.
+    """
     data_id = str(uuid.uuid4())
 
     if TESTING:
         data_id = "test-" + data_id
         session_id = "test-" + session_id
 
-    # Audio path in repo
     suffix = Path(audio).suffix
-    audio_p = f"{SPLIT}/audio/" + session_id + suffix
+    audio_repo_path = f"{SPLIT}/audio/{session_id}{suffix}"
 
-    if not hf_fs.exists(f"datasets/{DATASET_REPO}/{audio_p}"):
-        api.upload_file(
+    if not _fs.exists(f"datasets/{DATASET_REPO}/{audio_repo_path}"):
+        _api.upload_file(
             path_or_fileobj=str(audio),
-            path_in_repo=audio_p,
+            path_in_repo=audio_repo_path,
             repo_id=DATASET_REPO,
             repo_type="dataset",
         )
 
-    text = {
+    record = {
         "user_message": user_text,
         "model_response": model_response,
-        "file_name": "audio/" + session_id + suffix,  # has to be relative to metadata.jsonl
-        "original_fn": os.path.basename(audio),
+        "file_name": f"audio/{session_id}{suffix}",
+        "original_fn": Path(audio).name,
         "id": data_id,
         "session_id": session_id,
         "model_version": model_version,
     }
 
-    # Append to a jsonl file in the repo
-    # APPEND DOESN'T WORK, have to open first
-    if hf_fs.exists(f"datasets/{DATASET_REPO}/{SPLIT}/metadata.jsonl"):
-        with hf_fs.open(f"datasets/{DATASET_REPO}/{SPLIT}/metadata.jsonl", "r") as f:
+    line = json.dumps(record) + "\n"
+
+    # HfFileSystem doesn't support append, so read-then-write.
+    if _fs.exists(_METADATA_PATH):
+        with _fs.open(_METADATA_PATH, "r") as f:
             lines = f.readlines()
-        lines.append(json.dumps(text) + "\n")
-        with hf_fs.open(f"datasets/{DATASET_REPO}/{SPLIT}/metadata.jsonl", "w") as f:
+        lines.append(line)
+        with _fs.open(_METADATA_PATH, "w") as f:
             f.writelines(lines)
     else:
-        with hf_fs.open(f"datasets/{DATASET_REPO}/{SPLIT}/metadata.jsonl", "w") as f:
-            f.write(json.dumps(text) + "\n")
-
-    # Write a separate file instead
-    # with hf_fs.open(f"datasets/{DATASET_REPO}/{data_id}.json", "w") as f:
-    #     json.dump(text, f)
+        with _fs.open(_METADATA_PATH, "w") as f:
+            f.write(line)

infer.py

@@ -1,347 +0,0 @@
-# """Run NatureLM-audio over a set of audio files paths or a directory with audio files."""
-
-# import argparse
-# from pathlib import Path
-
-# import librosa
-# import numpy as np
-# import pandas as pd
-# import torch
-
-# from NatureLM.config import Config
-# from NatureLM.models import NatureLM
-# from NatureLM.processors import NatureLMAudioProcessor
-# from NatureLM.utils import move_to_device
-
-# _MAX_LENGTH_SECONDS = 10
-# _MIN_CHUNK_LENGTH_SECONDS = 0.5
-# _SAMPLE_RATE = 16000  # Assuming the model uses a sample rate of 16kHz
-# _AUDIO_FILE_EXTENSIONS = [".wav", ".mp3", ".flac", ".ogg", ".mp4"]  # Add other audio file formats as needed
-# _DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-# __root_dir = Path(__file__).parent.parent
-# _DEFAULT_CONFIG_PATH = __root_dir / "configs" / "inference.yml"
-
-
-# def load_model_and_config(
-#     cfg_path: str | Path = _DEFAULT_CONFIG_PATH, device: str = _DEVICE
-# ) -> tuple[NatureLM, Config]:
-#     """Load the NatureLM model and configuration.
-#     Returns:
-#         tuple: The loaded model and configuration.
-#     """
-#     model = NatureLM.from_pretrained("EarthSpeciesProject/NatureLM-audio")
-#     model = model.to(device).eval()
-#     model.llama_tokenizer.pad_token_id = model.llama_tokenizer.eos_token_id
-#     model.llama_model.generation_config.pad_token_id = model.llama_tokenizer.pad_token_id
-
-#     cfg = Config.from_sources(cfg_path)
-#     return model, cfg
-
-
-# def output_template(model_output: str, start_time: float, end_time: float) -> str:
-#     """Format the output of the model.
-
-#     Returns
-#     -------
-#     str
-#         Formatted string with timestamps and model output.
-#     """
-#     return f"#{start_time:.2f}s - {end_time:.2f}s#: {model_output}\n"
-
-
-# def sliding_window_inference(
-#     audio: str | Path | np.ndarray,
-#     query: str,
-#     processor: NatureLMAudioProcessor,
-#     model: NatureLM,
-#     cfg: Config,
-#     window_length_seconds: float = 10.0,
-#     hop_length_seconds: float = 10.0,
-#     input_sr: int = _SAMPLE_RATE,
-#     device: str = _DEVICE,
-# ) -> list[dict[str, any]]:
-#     """Run inference on a long audio file using sliding window approach.
-
-#     Args:
-#         audio (str | Path | np.ndarray): Path to the audio file.
-#         query (str): Query for the model.
-#         processor (NatureLMAudioProcessor): Audio processor.
-#         model (NatureLM): NatureLM model.
-#         cfg (Config): Model configuration.
-#         window_length_seconds (float): Length of the sliding window in seconds.
-#         hop_length_seconds (float): Hop length for the sliding window in seconds.
-#         input_sr (int): Sample rate of the audio file.
-
-#     Returns:
-#         str: The output of the model.
-
-#     Raises:
-#         ValueError: If the audio file is too short or if the audio file path is invalid.
-#     """
-#     if isinstance(audio, str) or isinstance(audio, Path):
-#         audio_array, input_sr = librosa.load(str(audio), sr=None, mono=False)
-#     elif isinstance(audio, np.ndarray):
-#         audio_array = audio
-#         print(f"Using provided sample rate: {input_sr}")
-
-#     audio_array = audio_array.squeeze()
-#     if audio_array.ndim > 1:
-#         axis_to_average = int(np.argmin(audio_array.shape))
-#         audio_array = audio_array.mean(axis=axis_to_average)
-#         audio_array = audio_array.squeeze()
-
-#     # Do initial check that the audio is long enough
-#     if audio_array.shape[-1] < int(_MIN_CHUNK_LENGTH_SECONDS * input_sr):
-#         raise ValueError(f"Audio is too short. Minimum length is {_MIN_CHUNK_LENGTH_SECONDS} seconds.")
-
-#     start = 0
-#     stride = int(hop_length_seconds * input_sr)
-#     window_length = int(window_length_seconds * input_sr)
-#     window_id = 0
-
-#     output = []  # Initialize output list
-#     while True:
-#         chunk = audio_array[start : start + window_length]
-#         if chunk.shape[-1] < int(_MIN_CHUNK_LENGTH_SECONDS * input_sr):
-#             break
-
-#         # Resamples, pads, truncates and creates torch Tensor
-#         audio_tensor, prompt_list = processor([chunk], [query], [input_sr])
-
-#         input_to_model = {
-#             "raw_wav": audio_tensor,
-#             "prompt": prompt_list[0],
-#             "audio_chunk_sizes": 1,
-#             "padding_mask": torch.zeros_like(audio_tensor).to(torch.bool),
-#         }
-#         input_to_model = move_to_device(input_to_model, device)
-
-#         # generate
-#         prediction: str = model.generate(input_to_model, cfg.generate, prompt_list)[0]
-
-#         # Post-process the prediction
-#         # prediction = output_template(prediction, start / input_sr, (start + window_length) / input_sr)
-#         # output += prediction
-#         output.append(
-#             {
-#                 "start_time": start / input_sr,
-#                 "end_time": (start + window_length) / input_sr,
-#                 "prediction": prediction,
-#                 "window_number": window_id,
-#             }
-#         )
-
-#         # Move the window
-#         start += stride
-
-#         if start + window_length > audio_array.shape[-1]:
-#             break
-
-#     return output
-
-
-# class Pipeline:
-#     """Pipeline for running NatureLM-audio inference on a list of audio files or audio arrays"""
-
-#     def __init__(self, model: NatureLM = None, cfg_path: str | Path = _DEFAULT_CONFIG_PATH) -> None:
-#         self.cfg_path = cfg_path
-
-#         # Load model and config
-#         if model is not None:
-#             self.cfg = Config.from_sources(cfg_path)
-#             self.model = model
-#         else:
-#             # Download model from hub
-#             self.model, self.cfg = load_model_and_config(cfg_path)
-
-#         self.processor = NatureLMAudioProcessor(sample_rate=_SAMPLE_RATE, max_length_seconds=_MAX_LENGTH_SECONDS)
-
-#     def __call__(
-#         self,
-#         audios: list[str | Path | np.ndarray],
-#         queries: str | list[str],
-#         window_length_seconds: float = 10.0,
-#         hop_length_seconds: float = 10.0,
-#         input_sample_rate: int = _SAMPLE_RATE,
-#         verbose: bool = False,
-#     ) -> list[str]:
-#         """Run inference on a list of audio file paths or a single audio file with a
-#         single query or a list of queries. If multiple queries are provided,
-#         we assume that they are in the same order as the audio files. If a single query
-#         is provided, it will be used for all audio files.
-
-#         Args:
-#             audios (list[str | Path | np.ndarray]): List of audio file paths or a single audio
-#                 file path or audio array(s)
-#             queries (str | list[str]): Queries for the model.
-#             window_length_seconds (float): Length of the sliding window in seconds. Defaults to 10.0.
-#             hop_length_seconds (float): Hop length for the sliding window in seconds. Defaults to 10.0.
-#             input_sample_rate (int): Sample rate of the audio. Defaults to 16000, which is the model's sample rate.
-#             verbose (bool): If True, print the output of the model for each audio file.
-#             Defaults to False.
-
-#         Returns:
-#             list[list[dict]]: List of model outputs for each audio file. Each output is a list of dictionaries
-#             containing the start time, end time, and prediction for each chunk of audio.
-
-#         Raises:
-#             ValueError: If the number of audio files and queries do not match.
-#         """
-#         if isinstance(audios, str) or isinstance(audios, Path):
-#             audios = [audios]
-
-#         if isinstance(queries, str):
-#             queries = [queries] * len(audios)
-
-#         if len(audios) != len(queries):
-#             raise ValueError("Number of audio files and queries must match.")
-
-#         # Run inference
-#         results = []
-#         for audio, query in zip(audios, queries, strict=False):
-#             output = sliding_window_inference(
-#                 audio,
-#                 query,
-#                 self.processor,
-#                 self.model,
-#                 self.cfg,
-#                 window_length_seconds,
-#                 hop_length_seconds,
-#                 input_sr=input_sample_rate,
-#             )
-#             results.append(output)
-#             if verbose:
-#                 print(f"Processed {audio}, model output:\n=======\n{output}\n=======")
-#         return results
-
-
-# def parse_args() -> argparse.Namespace:
-#     parser = argparse.ArgumentParser("Run NatureLM-audio inference")
-#     parser.add_argument(
-#         "-a",
-#         "--audio",
-#         type=str,
-#         required=True,
-#         help="Path to an audio file or a directory containing audio files",
-#     )
-#     parser.add_argument("-q", "--query", type=str, required=True, help="Query for the model")
-#     parser.add_argument(
-#         "--cfg-path",
-#         type=str,
-#         default="configs/inference.yml",
-#         help="Path to the configuration file for the model",
-#     )
-#     parser.add_argument(
-#         "--output_path",
-#         type=str,
-#         default="inference_output.jsonl",
-#         help="Output path for the results",
-#     )
-#     parser.add_argument(
-#         "--window_length_seconds",
-#         type=float,
-#         default=10.0,
-#         help="Length of the sliding window in seconds",
-#     )
-#     parser.add_argument(
-#         "--hop_length_seconds",
-#         type=float,
-#         default=10.0,
-#         help="Hop length for the sliding window in seconds",
-#     )
-#     args = parser.parse_args()
-
-#     return args
-
-
-# def main(
-#     cfg_path: str | Path,
-#     audio_path: str | Path,
-#     query: str,
-#     output_path: str,
-#     window_length_seconds: float,
-#     hop_length_seconds: float,
-# ) -> None:
-#     """Main function to run the NatureLM-audio inference script.
-#     It takes command line arguments for audio file path, query, output path,
-#     window length, and hop length. It processes the audio files and saves the
-#     results to a CSV file.
-
-#     Args:
-#         cfg_path (str | Path): Path to the configuration file.
-#         audio_path (str | Path): Path to the audio file or directory.
-#         query (str): Query for the model.
-#         output_path (str): Path to save the output results.
-#         window_length_seconds (float): Length of the sliding window in seconds.
-#         hop_length_seconds (float): Hop length for the sliding window in seconds.
-
-#     Raises:
-#         ValueError: If the audio file path is invalid or if the query is empty.
-#         ValueError: If no audio files are found.
-#         ValueError: If the audio file extension is not supported.
-#     """
-
-#     # Prepare sample
-#     audio_path = Path(audio_path)
-#     if audio_path.is_dir():
-#         audio_paths = []
-#         print(f"Searching for audio files in {str(audio_path)} with extensions {', '.join(_AUDIO_FILE_EXTENSIONS)}")
-#         for ext in _AUDIO_FILE_EXTENSIONS:
-#             audio_paths.extend(list(audio_path.rglob(f"*{ext}")))
-
-#         print(f"Found {len(audio_paths)} audio files in {str(audio_path)}")
-#     else:
-#         # check that the extension is valid
-#         if not any(audio_path.suffix == ext for ext in _AUDIO_FILE_EXTENSIONS):
-#             raise ValueError(
-#                 f"Invalid audio file extension. Supported extensions are: {', '.join(_AUDIO_FILE_EXTENSIONS)}"
-#             )
-#         audio_paths = [audio_path]
-
-#     # check that query is not empty
-#     if not query:
-#         raise ValueError("Query cannot be empty")
-#     if not audio_paths:
-#         raise ValueError("No audio files found. Please check the path or file extensions.")
-
-#     # Load model and config
-#     model, cfg = load_model_and_config(cfg_path)
-
-#     # Load audio processor
-#     processor = NatureLMAudioProcessor(sample_rate=_SAMPLE_RATE, max_length_seconds=_MAX_LENGTH_SECONDS)
-
-#     # Run inference
-#     results = {"audio_path": [], "output": []}
-#     for path in audio_paths:
-#         output = sliding_window_inference(
-#             path,
-#             query,
-#             processor,
-#             model,
-#             cfg,
-#             window_length_seconds,
-#             hop_length_seconds,
-#         )
-#         results["audio_path"].append(str(path))
-#         results["output"].append(output)
-#         print(f"Processed {path}, model output:\n=======\n{output}\n=======\n")
-
-#     # Save results as a csv
-#     output_path = Path(output_path)
-#     output_path.parent.mkdir(parents=True, exist_ok=True)
-
-#     df = pd.DataFrame(results)
-#     df.to_json(output_path, orient="records", lines=True)
-#     print(f"Results saved to {output_path}")
-
-
-# if __name__ == "__main__":
-#     args = parse_args()
-#     main(
-#         cfg_path=args.cfg_path,
-#         audio_path=args.audio,
-#         query=args.query,
-#         output_path=args.output_path,
-#         window_length_seconds=args.window_length_seconds,
-#         hop_length_seconds=args.hop_length_seconds,
-#     )

pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
     "torchaudio>=2.7.1",
     "matplotlib>=3.10.8",
     "numpy>=2.3.5",
+    "librosa>=0.9.2",
 ]
 
 [tool.uv.sources]