gagan/modelmerging_and_multiturn (#6)

- Model merging + audio cache (eabe9bf6849303d5a373ac512aa5ea120e210ea4)

NatureLM/config.py

@@ -136,6 +136,7 @@ class GenerateConfig(BaseModel, extra="forbid", validate_assignment=True):
     temperature: float
     repetition_penalty: float
     length_penalty: float
+    merging_alpha: float = 1.0
 
 
 class ModelConfig(BaseModel, extra="forbid", validate_assignment=True):

NatureLM/models/NatureLM.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import hashlib
 import logging
 import os
+from collections import OrderedDict
 from pathlib import Path
 from typing import Literal, Union
 
@@ -35,8 +37,98 @@ from .Qformer import BertConfig, BertLMHeadModel
 from .utils import StoppingCriteriaSub
 
 torch.backends.cuda.matmul.allow_tf32 = True
+auth_token = os.getenv("llama", None)
+
+
+class AudioEncodingCache:
+    """LRU cache for audio encoding with content-based hashing."""
+
+    def __init__(self, capacity: int = 100):
+        self.capacity = capacity
+        self.cache = OrderedDict()
+        self.hits = 0
+        self.misses = 0
+
+    def _compute_hash(
+        self, raw_wav: torch.Tensor, audio_padding_mask: torch.Tensor | None = None
+    ) -> str:
+        """Compute a hash key from the audio tensor and padding mask."""
+        # Use a sample of the tensor for efficiency (first, middle, last portions)
+        B, L = raw_wav.shape
+        sample_size = min(1000, L)  # Sample 1000 points or entire length if smaller
+
+        # Sample from beginning, middle, and end
+        indices = torch.cat(
+            [
+                torch.arange(min(sample_size // 3, L)),
+                torch.arange(L // 2, min(L // 2 + sample_size // 3, L)),
+                torch.arange(max(0, L - sample_size // 3), L),
+            ]
+        )
+
+        sampled_wav = raw_wav[:, indices].cpu().numpy().tobytes()
+
+        # Create hash from audio data, shape, and padding mask presence
+        hash_obj = hashlib.sha256(sampled_wav)
+        hash_obj.update(str(raw_wav.shape).encode())
+        hash_obj.update(str(raw_wav.dtype).encode())
+
+        if audio_padding_mask is not None:
+            mask_sample = audio_padding_mask[:, indices].cpu().numpy().tobytes()
+            hash_obj.update(mask_sample)
+            hash_obj.update(str(audio_padding_mask.shape).encode())
+        else:
+            hash_obj.update(b"no_mask")
+
+        return hash_obj.hexdigest()
+
+    def get(self, raw_wav: torch.Tensor, audio_padding_mask: torch.Tensor = None):
+        """Retrieve cached encoding if available."""
+        key = self._compute_hash(raw_wav, audio_padding_mask)
+
+        if key in self.cache:
+            self.hits += 1
+            # Move to end (most recently used)
+            self.cache.move_to_end(key)
+            return self.cache[key]
+
+        self.misses += 1
+        return None
+
+    def put(self, raw_wav: torch.Tensor, audio_padding_mask: torch.Tensor, value: tuple):
+        """Store encoding in cache (on CPU to save GPU memory)."""
+        key = self._compute_hash(raw_wav, audio_padding_mask)
+
+        # Move tensors to CPU for storage
+        audio_embeds, audio_atts = value
+        cached_value = (audio_embeds.cpu(), audio_atts.cpu())
+
+        # Add to cache
+        self.cache[key] = cached_value
+        self.cache.move_to_end(key)
+
+        # Evict oldest if over capacity
+        if len(self.cache) > self.capacity:
+            self.cache.popitem(last=False)
+
+    def clear(self):
+        """Clear the cache."""
+        self.cache.clear()
+        self.hits = 0
+        self.misses = 0
+
+    def get_stats(self):
+        """Get cache statistics."""
+        total = self.hits + self.misses
+        hit_rate = self.hits / total if total > 0 else 0
+        return {
+            "hits": self.hits,
+            "misses": self.misses,
+            "hit_rate": hit_rate,
+            "size": len(self.cache),
+            "capacity": self.capacity,
+        }
 
-auth_token = os.getenv('llama')
 
 class NatureLM(nn.Module, PyTorchModelHubMixin):
     def __init__(
@@ -65,9 +157,16 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
         max_txt_len: int = 128,
         end_sym: str = "</s>",
         device: str = "cuda",
+        audio_encoding_cache_size: int = 100,
     ):
         super().__init__()
 
+        self.audio_encoding_cache = (
+            AudioEncodingCache(capacity=audio_encoding_cache_size)
+            if audio_encoding_cache_size > 0
+            else None
+        )
+
         self.beats_path = beats_path
         self.beats_cfg = beats_cfg
         self.use_audio_Qformer = use_audio_Qformer
@@ -84,7 +183,9 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
 
         logging.info(f"Llama path: {llama_path}")
         logging.info("Loading Llama Tokenizer")
-        self.llama_tokenizer = AutoTokenizer.from_pretrained(llama_path, use_fast=False, use_auth_token=auth_token)
+        self.llama_tokenizer = AutoTokenizer.from_pretrained(
+            llama_path, use_fast=False, use_auth_token=auth_token
+        )
         self.llama_tokenizer.add_special_tokens({"pad_token": "[PAD]"})
         self.llama_tokenizer.padding_side = "right"
 
@@ -95,7 +196,6 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
                 torch_dtype=torch.float32,
                 attn_implementation="eager",
                 device_map="cpu",
-                use_auth_token=auth_token
             )
             # An issue with tiny-llama is that pad_token_id was set to -1, but
             # model.save_pretrained checks generation configs and does not allow -1 as
@@ -106,7 +206,6 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
                 llama_path,
                 torch_dtype=torch.bfloat16,
                 attn_implementation=flash_attn,
-                use_auth_token=auth_token
             )
 
         self.llama_model.resize_token_embeddings(len(self.llama_tokenizer))
@@ -135,7 +234,9 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
         self.beats = BEATs(cfg=BEATsConfig(dict(self.beats_cfg)))
 
         if self.beats_path:
-            beats_ckpt = universal_torch_load(self.beats_path, cache_mode="none", map_location="cpu")
+            beats_ckpt = universal_torch_load(
+                self.beats_path, cache_mode="none", map_location="cpu"
+            )
             self.beats.load_state_dict(beats_ckpt["model"])
 
         self.ln_audio = nn.LayerNorm(self.beats.cfg.encoder_embed_dim)
@@ -336,11 +437,15 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
             audio_embeds = self.ln_audio(audio_embeds)
 
             # Generate attention mask
-            audio_atts = torch.ones(audio_embeds.size()[:-1], dtype=torch.long).to(audio_embeds.device)
+            audio_atts = torch.ones(audio_embeds.size()[:-1], dtype=torch.long).to(
+                audio_embeds.device
+            )
 
             if self.window_level_Qformer:
                 B, T, C = audio_embeds.shape  # batch, T, Channels
-                kernel = round(1500 * self.second_per_window / 30.0)  # 160 ms patches; calculate kernel size
+                kernel = round(
+                    1500 * self.second_per_window / 30.0
+                )  # 160 ms patches; calculate kernel size
                 stride = round(1500 * self.second_stride / 30.0)  # Calculate stride size
                 kernel = (1, kernel)
                 stride = (1, stride)
@@ -360,7 +465,9 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
                     audio_embeds_overlap, [0, 3, 2, 1]
                 )  # (B, num_windows, kernel_size, C)
                 audio_embeds = audio_embeds_overlap.reshape(-1, kernel[1], C)
-                audio_atts = torch.ones(audio_embeds.size()[:-1], dtype=torch.long).to(audio_embeds.device)
+                audio_atts = torch.ones(audio_embeds.size()[:-1], dtype=torch.long).to(
+                    audio_embeds.device
+                )
 
                 # Q-Former mechanism
                 query_tokens = self.audio_query_tokens.expand(audio_embeds.shape[0], -1, -1)
@@ -376,13 +483,19 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
                 if self.window_level_Qformer:
                     audio_embeds = audio_embeds.view(B, -1, audio_embeds.size(2)).contiguous()
 
-            audio_atts = torch.ones(audio_embeds.size()[:-1], dtype=torch.long).to(audio_embeds.device)
+            audio_atts = torch.ones(audio_embeds.size()[:-1], dtype=torch.long).to(
+                audio_embeds.device
+            )
 
         elif self.htsat:
             # HTSAT processing
             audio_embeds = self.ln_audio(audio_embeds)
-            audio_embeds = self.audio_llama_proj(audio_embeds).reshape(-1, 30, self.llama_model.config.hidden_size)
-            audio_atts = torch.ones(audio_embeds.size()[:-1], dtype=torch.long).to(audio_embeds.device)
+            audio_embeds = self.audio_llama_proj(audio_embeds).reshape(
+                -1, 30, self.llama_model.config.hidden_size
+            )
+            audio_atts = torch.ones(audio_embeds.size()[:-1], dtype=torch.long).to(
+                audio_embeds.device
+            )
 
         else:
             raise NotImplementedError("no audio qformer or max pooling")
@@ -390,9 +503,32 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
         return audio_embeds, audio_atts
 
     def encode_audio(self, raw_wav, audio_padding_mask=None):
+        # Only use cache during inference (not training)
+        if self.audio_encoding_cache is not None and not self.training:
+            cached_result = self.audio_encoding_cache.get(raw_wav, audio_padding_mask)
+            if cached_result is not None:
+                print("#### Audio encoding cache hit ####")
+                # Move cached tensors back to the model's device
+                audio_embeds, audio_atts = cached_result
+                return audio_embeds.to(self.device), audio_atts.to(self.device)
+
+        # Compute encoding if not cached
         with torch.autocast(self.device.type, dtype=torch.bfloat16):
             audio_embeds, audio_pad_mask = self.beats(raw_wav, padding_mask=audio_padding_mask)
-            return self._encode_auditory_feature(audio_embeds=audio_embeds, audio_pad_mask=audio_pad_mask)
+            result = self._encode_auditory_feature(
+                audio_embeds=audio_embeds, audio_pad_mask=audio_pad_mask
+            )
+
+        # Store in cache if enabled and in inference mode
+        if self.audio_encoding_cache is not None and not self.training:
+            self.audio_encoding_cache.put(raw_wav, audio_padding_mask, result)
+
+        return result
+
+    def clear_audio_embed_cache(self):
+        """Clear the audio encoding cache."""
+        if self.audio_encoding_cache is not None:
+            self.audio_encoding_cache.clear()
 
     def prompt_wrap(self, audio_embeds, audio_atts, prompt: list[str]):
         """Merge audio embeddings with embeddings of the tokens in the prompt.
@@ -440,7 +576,9 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
             wrapped_atts = []
 
             for part in prompt_parts:
-                tokens = self.llama_tokenizer(part, return_tensors="pt", add_special_tokens=False).to(device)
+                tokens = self.llama_tokenizer(
+                    part, return_tensors="pt", add_special_tokens=False
+                ).to(device)
                 part_embeds = self.llama_embed_tokens(tokens.input_ids).squeeze(0)
                 part_atts = tokens.attention_mask.squeeze(0)
                 wrapped_embeds.append(part_embeds)
@@ -506,7 +644,9 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
 
         # BOS token embeddings
         bos_token_id = self.llama_tokenizer.bos_token_id
-        bos = torch.full((batch_size, 1), bos_token_id, dtype=torch.long, device=audio_embeds.device)
+        bos = torch.full(
+            (batch_size, 1), bos_token_id, dtype=torch.long, device=audio_embeds.device
+        )
         bos_embeds = self.llama_embed_tokens(bos)
 
         # Prepare lists to collect per-sample embeddings, attention masks, and targets
@@ -521,7 +661,9 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
 
             # Extract non-padded text embeddings and attention mask
             text_embed = to_regress_embeds[i][to_regress_tokens.attention_mask[i].bool()]
-            text_att = to_regress_tokens.attention_mask[i][to_regress_tokens.attention_mask[i].bool()]
+            text_att = to_regress_tokens.attention_mask[i][
+                to_regress_tokens.attention_mask[i].bool()
+            ]
 
             # Extract corresponding targets for the text tokens
             target = targets[i][to_regress_tokens.attention_mask[i].bool()]
@@ -581,7 +723,9 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
             shift_logits.view(-1, nvocab),  # Flatten to [batch_size * (seq_len-1), vocab_size]
             shift_labels.view(-1),  # Flatten to [batch_size * (seq_len-1)]
         )
-        loss_per_token = loss_per_token.view(shift_labels.size())  # Reshape back to [batch_size, seq_len-1]
+        loss_per_token = loss_per_token.view(
+            shift_labels.size()
+        )  # Reshape back to [batch_size, seq_len-1]
 
         # Create mask
         mask = shift_labels != -100  # [batch_size, seq_len-1]
@@ -597,7 +741,9 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
             predicted_tokens = shift_logits.argmax(dim=-1)  # [batch_size, seq_len-1]
 
             # Compute per-example correct counts
-            correct_per_sample = ((predicted_tokens == shift_labels) & mask).sum(dim=1).float()  # [batch_size]
+            correct_per_sample = (
+                ((predicted_tokens == shift_labels) & mask).sum(dim=1).float()
+            )  # [batch_size]
             total_tokens_per_sample = mask.sum(dim=1).float()  # [batch_size]
 
             # Total correct and total tokens across the batch
@@ -615,8 +761,31 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
 
         return {"loss": loss, "per_example_loss": loss_per_example}
 
+    def model_merging_scaling(self, merging_alpha, adapter_name="default"):
+        """
+        Performs model merging with the base model by adjusting the scaling of the LoRA adapters as described in
+        "Model Merging Improves Zero-Shot Generalization in Bioacoustic Foundation Models"
+        (https://arxiv.org/abs/2511.05171).
+
+        The best value for alpha is task- and dataset-specific, but the paper found alpha values between
+        0.4 and 0.6 to perform generally well.
+
+        Args:
+            merging_alpha: The merging_alpha used for interpolation.
+            adapter_name (str): The name of the adapter to rescale when merging.
+        """
+
+        for module in self.llama_model.modules():
+            # Check if the module is a LoRA layer and has the specified adapter
+            if hasattr(module, "r") and isinstance(module.r, dict) and adapter_name in module.r:
+                module.scaling[adapter_name] = merging_alpha * module.scaling[adapter_name]
+
     @torch.inference_mode()
-    def generate(self, samples, generate_cfg, prompts):
+    def generate(self, samples, generate_cfg, prompts) -> list[str]:
+        merging_alpha = getattr(generate_cfg, "merging_alpha", 1.0)
+        if merging_alpha != 1.0:
+            self.model_merging_scaling(merging_alpha)
+
         batch_size = len(prompts)
 
         raw_wav = samples["raw_wav"]
@@ -645,7 +814,7 @@ class NatureLM(nn.Module, PyTorchModelHubMixin):
         stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
 
         with torch.autocast(self.device.type, dtype=torch.bfloat16):
-            outputs = self.llama_model.generate(  
+            outputs = self.llama_model.generate(  # TODO: Wrap the llama_model with outlines https://outlines-dev.github.io/outlines/reference/models/transformers/
                 inputs_embeds=embeds.bfloat16(),
                 max_new_tokens=generate_cfg.max_new_tokens,
                 stopping_criteria=stopping_criteria,

app.py

@@ -1,4 +1,5 @@
 import spaces
+import uuid
 import warnings
 import traceback
 import numpy as np
@@ -17,19 +18,29 @@ from NatureLM.infer import Pipeline
 
 from data_store import upload_data
 
+
 warnings.filterwarnings("ignore")
 SAMPLE_RATE = 16000  # Default sample rate for NatureLM-audio
 DEVICE: str = "cuda" if torch.cuda.is_available() else "cpu"
+MIN_AUDIO_DURATION: float = 0.5  # seconds
+MAX_HISTORY_TURNS = (
+    3  # Maximum number of conversation turns to include in context (user + assistant pairs)
+)
 
 # Load model at startup if CUDA is available
 print(f"Device: {DEVICE}")
-if DEVICE == "cuda":
-    print("CUDA available, loading model at startup...")
-    model = NatureLM.from_pretrained("EarthSpeciesProject/NatureLM-audio")
-    model = model.eval().to(DEVICE)
-    model = Pipeline(model)
-else:
-    print("CUDA not available, model will not be loaded at startup")
+model = NatureLM.from_pretrained("EarthSpeciesProject/NatureLM-audio")
+model = model.eval().to(DEVICE)
+model = Pipeline(model)
+
+
+def check_audio_duration_greater(audio_path: str) -> bool:
+    """Check the duration of the audio file."""
+    info = torchaudio.info(audio_path)
+    duration = info.num_frames / info.sample_rate
+    if not duration >= MIN_AUDIO_DURATION:
+        raise gr.Error(f"Audio duration must be at least {MIN_AUDIO_DURATION} seconds.")
+
 
 def get_spectrogram(audio: torch.Tensor) -> plt.Figure:
     """Generate a spectrogram from the audio tensor."""
@@ -98,13 +109,11 @@ def prompt_lm(
     hop_length_seconds: float = 10.0,
 ) -> list[str]:
     """Generate response using the model
-
     Args:
         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[str]: List of generated responses for each audio-query pair
     """
@@ -157,28 +166,61 @@ def add_user_query(chatbot_history: list[dict], chat_input: str) -> list[dict]:
     return chatbot_history
 
 
-def send_data_to_hub(chatbot_history: list[dict], audio: str):
+def send_data_to_hub(chatbot_history: list[dict], audio: str, session_id: str):
     """Upload data to hub"""
     if not chatbot_history or len(chatbot_history) < 2:
         return
     user_text = chatbot_history[-2]["content"]
     model_response = chatbot_history[-1]["content"]
-    upload_data(audio, user_text, model_response)
+    upload_data(audio, user_text, model_response, session_id)
 
 
 def get_response(chatbot_history: list[dict], audio_input: str) -> list[dict]:
-    """Generate response from the model based on user input and audio file"""
+    """Generate response from the model based on user input and audio file with conversation history"""
     try:
-        # Get the last user message from chat history
+        # Warn if conversation is getting long
+        num_turns = len(chatbot_history)
+        if num_turns > MAX_HISTORY_TURNS * 2:  # Each turn = user + assistant message
+            gr.Warning(
+                "⚠️ Long conversations may affect response quality. 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
-        print("\nUser message:", last_user_message)
+
+        # 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
+
+            full_prompt = (
+                "Previous conversation:\n"
+                + "\n".join(recent_context)
+                + "\n\nCurrent question: "
+                + last_user_message
+            )
+        else:
+            full_prompt = last_user_message
+
+        print("\nFull prompt with history:", full_prompt)
+
         response = prompt_lm(
             audios=[audio_input],
-            queries=[last_user_message.strip()],
+            queries=[full_prompt.strip()],
             window_length_seconds=100_000,
             hop_length_seconds=100_000,
         )
@@ -192,7 +234,7 @@ def get_response(chatbot_history: list[dict], audio_input: str) -> list[dict]:
         print(f"Error generating response: {e}")
         traceback.print_exc()
         response = "Error generating response. Please try again."
-        
+
     # Add model response to chat history
     chatbot_history.append({"role": "assistant", "content": response})
 
@@ -201,17 +243,7 @@ def get_response(chatbot_history: list[dict], audio_input: str) -> list[dict]:
 
 def main(
     assets_dir: Path,
-    cfg_path: str | Path,
-    options: list[str] = [],
 ):
-    # Load configuration
-    try:
-        cfg = Config.from_sources(yaml_file=cfg_path, cli_args=options)
-        print("Configuration loaded successfully")
-    except Exception as e:
-        print(f"Warning: Could not load config: {e}")
-        print("Running in demo mode")
-
     # Check if assets directory exists, if not create a placeholder
     if not assets_dir.exists():
         print(f"Warning: Assets directory {assets_dir} does not exist")
@@ -248,13 +280,11 @@ def main(
         "Caption the audio (Humpback Whale)": [str(whale_audio), "Caption the audio."],
     }
 
-    gr.set_static_paths(paths=[Path.cwd().absolute()/"assets"])
+    gr.set_static_paths(paths=[Path.cwd().absolute() / "assets"])
 
     with gr.Blocks(
         title="NatureLM-audio",
-        theme=gr.themes.Base(
-            primary_hue="blue", font=[gr.themes.GoogleFont("Noto Sans")]
-        )
+        theme=gr.themes.Base(primary_hue="blue", font=[gr.themes.GoogleFont("Noto Sans")]),
     ) as app:
         with gr.Row():
             gr.HTML("""
@@ -272,7 +302,8 @@ def main(
 
         with gr.Tabs():
             with gr.Tab("Analyze Audio"):
-                uploaded_audio = gr.State()
+                session_id = gr.State(str(uuid.uuid4()))
+                # uploaded_audio = gr.State()
                 # Status indicator
                 # status_text = gr.Textbox(
                 #     value=model_manager.get_status(),
@@ -297,8 +328,6 @@ def main(
                     """,
                         padding=False,
                     )
-                    
-
 
                 with gr.Column(visible=True) as upload_section:
                     audio_input = gr.Audio(
@@ -307,6 +336,14 @@ def main(
                         interactive=True,
                         sources=["upload"],
                     )
+                    # check that audio duration is greater than MIN_AUDIO_DURATION
+                    # raise
+                    audio_input.change(
+                        fn=check_audio_duration_greater,
+                        inputs=[audio_input],
+                        outputs=[],
+                    )
+
                 with gr.Accordion(
                     label="Toggle Spectrogram", open=False, visible=False
                 ) as spectrogram:
@@ -445,13 +482,11 @@ def main(
                     [clear_button],
                 ).then(
                     send_data_to_hub,
-                    [chatbot, audio_input],
+                    [chatbot, audio_input, session_id],
                     None,
                 )
 
-                clear_button.click(
-                    lambda: gr.ClearButton(visible=False), None, [clear_button]
-                )
+                clear_button.click(lambda: gr.ClearButton(visible=False), None, [clear_button])
 
             with gr.Tab("Sample Library"):
                 with gr.Row():
@@ -494,7 +529,6 @@ def main(
                                             type="filepath",
                                             show_download_button=True,
                                         )
-                        
 
             with gr.Tab("💡 Help"):
                 gr.HTML("""  
@@ -519,7 +553,6 @@ def main(
                             </ol> 
                         <p></p>
                         </div>      
-
                         <div class="guide-section">   
                             <h3>Tips</h3>   
                                 <b>Prompting Best Practices</b>
@@ -561,7 +594,6 @@ def main(
                     background: white;   
                     flex: 1; 
                 }
-
                 #chat-input .submit-button {  
                     padding: 10px;
                     margin: 2px 6px;
@@ -590,7 +622,6 @@ def main(
                     color: #374151; 
                     margin-bottom: 4px;
                 }
-
                 .banner .banner-text {
                     style="font-size: 14px; 
                     color: #6b7280; 
@@ -609,7 +640,6 @@ def main(
                     display: inline-block;
                     transition: background 0.2s ease;
                 }
-
                 .link-btn:hover { 
                     background: #2563eb; 
                 } 
@@ -635,12 +665,10 @@ def main(
                     #chat-input {
                         background: #1e1e1e;  
                     }
-
                     #chat-input textarea {
                         background: #1e1e1e;
                         color: white;  
                     }
-
                     .banner {  
                         background: #1e1e1e;  
                         color: white;
@@ -657,8 +685,6 @@ def main(
 # Create and launch the app
 app = main(
     assets_dir=Path("assets"),
-    cfg_path=Path("configs/inference.yml"),
-    options=[],
 )
 
 if __name__ == "__main__":

configs/inference.yml

@@ -59,3 +59,4 @@ generate:
   temperature: 0.1
   repetition_penalty: 1.0
   length_penalty: 1.0
+  merging_alpha: 0.5

data_store.py

@@ -7,33 +7,38 @@ 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))
+api = HfApi(token=os.getenv("HF_TOKEN", None))
 # Upload audio
 # check if file exists
-hf_fs = HfFileSystem(token=os.getenv("HF_TOKEN",None))
+hf_fs = HfFileSystem(token=os.getenv("HF_TOKEN", None))
 
 
-def upload_data(audio: str | Path, user_text: str, model_response: str):
+def upload_data(audio: str | Path, user_text: str, model_response: str, session_id: str = ""):
     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/" + data_id + suffix
+    audio_p = f"{SPLIT}/audio/" + session_id + suffix
 
-    api.upload_file(
-        path_or_fileobj=str(audio),
-        path_in_repo=audio_p,
-        repo_id=DATASET_REPO,
-        repo_type="dataset",
-    )
+    if not hf_fs.exists(f"datasets/{DATASET_REPO}/{audio_p}"):
+        api.upload_file(
+            path_or_fileobj=str(audio),
+            path_in_repo=audio_p,
+            repo_id=DATASET_REPO,
+            repo_type="dataset",
+        )
 
     text = {
         "user_message": user_text,
         "model_response": model_response,
-        "file_name": "audio/" + data_id + suffix, # has to be relative to metadata.jsonl
+        "file_name": "audio/" + session_id + suffix,  # has to be relative to metadata.jsonl
         "original_fn": os.path.basename(audio),
         "id": data_id,
+        "session_id": session_id,
     }
 
     # Append to a jsonl file in the repo