|
|
from typing import Dict, List, Any |
|
|
from NatureLM.models import NatureLM |
|
|
from NatureLM.infer import Pipeline |
|
|
import numpy as np |
|
|
import torch |
|
|
import os |
|
|
|
|
|
class EndpointHandler(): |
|
|
def __init__(self, path=""): |
|
|
auth_token = os.environ.get("LLAMA_TOK") |
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
model = NatureLM.from_pretrained("EarthSpeciesProject/NatureLM-audio", device=device, hf_auth_token=auth_token) |
|
|
self.model = model.eval().to(device) |
|
|
|
|
|
cfg_path = "inference.yml" |
|
|
if not os.path.exists(cfg_path): |
|
|
script_dir = os.path.dirname(os.path.abspath(__file__)) |
|
|
cfg_path = os.path.join(script_dir, "inference.yml") |
|
|
|
|
|
if not os.path.exists(cfg_path): |
|
|
raise FileNotFoundError(f"inference.yml not found at {cfg_path}. Current directory contents: {os.listdir('.')}") |
|
|
|
|
|
self.pipeline = Pipeline(model=self.model, cfg_path=cfg_path) |
|
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
|
""" |
|
|
Process audio list of floats with NatureLM-audio model. |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
data : Dict[str, Any] |
|
|
Dictionary containing: |
|
|
- inputs : list[float] |
|
|
Audio data as list of floats |
|
|
- query : str |
|
|
Question to ask about the audio |
|
|
- sample_rate : int, optional |
|
|
Audio sample rate, default 16000 |
|
|
|
|
|
Returns |
|
|
------- |
|
|
List[Dict[str, Any]] |
|
|
List containing result dictionary with 'result' and 'query' keys, |
|
|
or error dictionary with 'error' key |
|
|
""" |
|
|
|
|
|
audio = data.get("inputs") |
|
|
query = data.get("query", "") |
|
|
sample_rate = data.get("sample_rate", 16000) |
|
|
|
|
|
if audio is None: |
|
|
return [{"error": "No audio data provided"}] |
|
|
|
|
|
if not query: |
|
|
return [{"error": "No query provided"}] |
|
|
|
|
|
|
|
|
if isinstance(audio, list): |
|
|
audio = np.array(audio, dtype=np.float32) |
|
|
elif not isinstance(audio, np.ndarray): |
|
|
return [{"error": f"Audio data must be a list or numpy array, got {type(audio)}"}] |
|
|
|
|
|
try: |
|
|
|
|
|
results = self.pipeline( |
|
|
audios=[audio], |
|
|
queries=query, |
|
|
input_sample_rate=sample_rate |
|
|
) |
|
|
|
|
|
return [{"result": results, "query": query}] |
|
|
|
|
|
except Exception as e: |
|
|
return [{"error": f"Error processing audio: {str(e)}"}] |
