Just DiVA

app.py

@@ -1,259 +1,82 @@
 import copy
 import os
 import random
-import shutil
 import sys
-from pathlib import Path
 
+import xxhash
 import gradio as gr
 import librosa
 import numpy as np
 import soundfile as sf
-import spaces
 import torch
 import torch.nn.functional as F
 from accelerate import infer_auto_device_map
 from datasets import Audio
-from huggingface_hub import CommitScheduler, delete_file, hf_hub_download
+from models.salmonn import SALMONN
 from safetensors.torch import load, load_model
-from tinydb import TinyDB
+import spaces
 from torch import nn
 from transformers import (
-    AutoModel,
     AutoModelForCausalLM,
     AutoProcessor,
     AutoTokenizer,
     LlamaForCausalLM,
     TextIteratorStreamer,
     WhisperForConditionalGeneration,
+    AutoProcessor,
+    AutoModel
 )
 from transformers.generation import GenerationConfig
-import spaces
-
-# Set an environment variable
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
-model_id = "meta-llama/Meta-Llama-3-8B"
-# Load the tokenizer and model
-AutoTokenizer.from_pretrained(model_id)
-AutoModelForCausalLM.from_pretrained(model_id)
-
-#from models.salmonn import SALMONN
-
-DB_PATH = "user_study.json"
-DB_DATASET_ID = "WillHeld/DiVAVotes"
-
-# Download existing DB
-if not os.path.isfile(DB_PATH):
-    print("Downloading DB...")
-    try:
-        cache_path = hf_hub_download(
-            repo_id=DB_DATASET_ID, repo_type="dataset", filename=DB_PATH
-        )
-        shutil.copyfile(cache_path, DB_PATH)
-        print("Downloaded DB")
-    except Exception as e:
-        print("Error while downloading DB:", e)
-
-db = TinyDB(DB_PATH)
-
-# Sync local DB with remote repo every 5 minute (only if a change is detected)
-scheduler = CommitScheduler(
-    repo_id=DB_DATASET_ID,
-    repo_type="dataset",
-    folder_path=Path(DB_PATH).parent,
-    every=5,
-    allow_patterns=DB_PATH,
-)
-
-tokenizer = AutoTokenizer.from_pretrained("WillHeld/via-llama")
-prefix = torch.tensor([128000, 128006, 882, 128007, 271]).to("cuda")
-pre_user_suffix = torch.tensor([271]).to("cuda")
-final_header = torch.tensor([128009, 128006, 78191, 128007, 271]).to("cuda")
-cache = None
 anonymous = False
 
-resampler = Audio(sampling_rate=16_000)
-
-
-qwen_tokenizer = AutoTokenizer.from_pretrained(
-    "Qwen/Qwen-Audio-Chat", trust_remote_code=True
-)
-qwen_model = AutoModelForCausalLM.from_pretrained(
-    "Qwen/Qwen-Audio-Chat",
-    device_map="auto",
-    trust_remote_code=True,
-    torch_dtype=torch.float16,
-).eval()
-
-qwen_model.generation_config = GenerationConfig.from_pretrained(
-    "Qwen/Qwen-Audio-Chat",
-    trust_remote_code=True,
-    do_sample=False,
-    top_k=50,
-    top_p=1.0,
-)
-
-
-# salmonn_model = SALMONN(
-#     ckpt="./SALMONN_PATHS/salmonn_v1.pth",
-#     whisper_path="./SALMONN_PATHS/whisper-large-v2",
-#     beats_path="./SALMONN_PATHS/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt2.pt",
-#     vicuna_path="./SALMONN_PATHS/vicuna-13b-v1.1",
-#     low_resource=False,
-#     device="cuda",
-# )
-# salmonn_tokenizer = salmonn_model.llama_tokenizer
-
-
-diva = AutoModel.from_pretrained("WillHeld/DiVA-llama-3-v0-8b", trust_remote_code=True, speech_encoder_device="cuda")
-
-# @spaces.GPU
-# @torch.no_grad
-# def salmonn_fwd(audio_input, prompt, do_sample=False, temperature=0.001):
-#     if audio_input == None:
-#         return ""
-#     sr, y = audio_input
-#     y = y.astype(np.float32)
-#     y /= np.max(np.abs(y))
-#     a = resampler.decode_example(
-#         resampler.encode_example({"array": y, "sampling_rate": sr})
-#     )
-#     sf.write("tmp.wav", a["array"], a["sampling_rate"], format="wav")
-#     streamer = TextIteratorStreamer(salmonn_tokenizer)
-#     with torch.cuda.amp.autocast(dtype=torch.float16):
-#         llm_message = salmonn_model.generate(
-#             wav_path="tmp.wav",
-#             prompt=prompt,
-#             do_sample=False,
-#             top_p=1.0,
-#             temperature=0.0,
-#             device="cuda:0",
-#             streamer=streamer,
-#         )
-
-#     response = ""
-#     for new_tokens in streamer:
-#         response += new_tokens
-#         yield response.replace("</s>", "")
+diva_model = AutoModel.from_pretrained("WillHeld/DiVA-llama-3-v0-8b", trust_remote_code=True)
 
+resampler = Audio(sampling_rate=16_000)
 
 @spaces.GPU
 @torch.no_grad
-def qwen_audio(audio_input, prompt, do_sample=False, temperature=0.001):
-    if audio_input == None:
-        return ""
+def diva_audio(audio_input, do_sample=False, temperature=0.001):
     sr, y = audio_input
+    x = xxhash.xxh32(bytes(y)).hexdigest()
     y = y.astype(np.float32)
     y /= np.max(np.abs(y))
     a = resampler.decode_example(
         resampler.encode_example({"array": y, "sampling_rate": sr})
     )
-    sf.write("tmp.wav", a["array"], a["sampling_rate"], format="wav")
-    query = qwen_tokenizer.from_list_format([{"audio": "tmp.wav"}, {"text": prompt}])
-
-    response, history = qwen_model.chat(
-        qwen_tokenizer,
-        query=query,
-        system="You are a helpful assistant.",
-        history=None,
-    )
-    return response
-
+    yield from diva_model.generate_stream(a["array"], None, do_sample=do_sample, max_new_tokens = 256)
 
 @spaces.GPU
-@torch.no_grad
-def via(audio_input, prompt, do_sample=False, temperature=0.001):
-    if audio_input == None:
-        return ""
-    sr, y = audio_input
-    y = y.astype(np.float32)
-    y /= np.max(np.abs(y))
-    a = resampler.decode_example(
-        resampler.encode_example({"array": y, "sampling_rate": sr})
-    )
-
-    audio = a["array"]
-
-    yield from diva.generate_stream(audio, prompt)
-
-
-def transcribe(audio_input, text_prompt, state, model_order):
-    yield (
-        gr.Button(
-            value="Waiting in queue for GPU time...",
-            interactive=False,
-            variant="primary",
-        ),
-        "",
-        "",
-        "",
-        gr.Button(visible=False),
-        gr.Button(visible=False),
-        gr.Button(visible=False),
-        state,
-    )
+def transcribe(audio_input, state, model_order):
     if audio_input == None:
         return (
             "",
-            "",
-            "",
-            gr.Button(visible=False),
-            gr.Button(visible=False),
-            gr.Button(visible=False),
             state,
         )
 
-    def gen_from_via():
-        via_resp = via(audio_input, text_prompt)
-        for resp in via_resp:
-            v_resp = gr.Textbox(
+    def gen_from_diva():
+        diva_resp = diva_audio(audio_input)
+        for resp in  diva_resp:
+            d_resp = gr.Textbox(
                 value=resp,
                 visible=True,
                 label=model_names[0] if not anonymous else f"Model {order}",
             )
-            yield (v_resp, s_resp, q_resp)
-
-    # def gen_from_salmonn():
-    #     salmonn_resp = salmonn_fwd(audio_input, text_prompt)
-    #     for resp in salmonn_resp:
-    #         s_resp = gr.Textbox(
-    #             value=resp,
-    #             visible=True,
-    #             label=model_names[1] if not anonymous else f"Model {order}",
-    #         )
-    #         yield (v_resp, s_resp, q_resp)
-
-    def gen_from_qwen():
-        qwen_resp = qwen_audio(audio_input, text_prompt)
-        q_resp = gr.Textbox(
-            value=qwen_resp,
-            visible=True,
-            label=model_names[2] if not anonymous else f"Model {order}",
-        )
-        yield (v_resp, s_resp, q_resp)
+            yield (d_resp)
 
     spinner_id = 0
     spinners = ["◐ ", "◓ ", "◑", "◒"]
-    initial_responses = [("", "", "")]
+    initial_responses = [("")]
     resp_generators = [
-        gen_from_via(),
-        # gen_from_salmonn(),
-        gen_from_qwen(),
+        gen_from_diva(),
     ]
     order = -1
-    resp_generators = [
-        resp_generators[model_order[0]],
-        #resp_generators[model_order[1]],
-        resp_generators[model_order[1]],
-    ]
+    resp_generators = [resp_generators[model_order[0]]]
     for generator in [initial_responses, *resp_generators]:
         order += 1
         for resps in generator:
-            v_resp, s_resp, q_resp = resps
+            s_resp, q_resp = resps
             resp_1 = resps[model_order[0]]
-            resp_2 = s_resp #resps[model_order[1]]
-            resp_3 = resps[model_order[1]]
             spinner = spinners[spinner_id]
             spinner_id = (spinner_id + 1) % 4
             yield (
@@ -263,11 +86,6 @@ def transcribe(audio_input, text_prompt, state, model_order):
                     variant="primary",
                 ),
                 resp_1,
-                resp_2,
-                resp_3,
-                gr.Button(visible=False),
-                gr.Button(visible=False),
-                gr.Button(visible=False),
                 state,
             )
     yield (
@@ -276,10 +94,6 @@ def transcribe(audio_input, text_prompt, state, model_order):
         ),
         resp_1,
         resp_2,
-        resp_3,
-        gr.Button(visible=True),
-        gr.Button(visible=False),
-        gr.Button(visible=True),
         responses_complete(state),
     )
 
@@ -287,7 +101,7 @@ def transcribe(audio_input, text_prompt, state, model_order):
 def on_page_load(state, model_order):
     if state == 0:
         gr.Info(
-            "Record what you want to say to your AI Assistant! All Audio recordings are stored only temporarily and will be erased as soon as you exit this page."
+            "Record something you'd say to an AI Assistant! Think about what you usually use Siri, Google Assistant, or ChatGPT for."
         )
         state = 1
         if anonymous:
@@ -298,53 +112,27 @@ def on_page_load(state, model_order):
 def recording_complete(state):
     if state == 1:
         gr.Info(
-            "Submit your recording to get responses from all three models! You can also influence the model responses with an optional prompt."
+            "Once you submit your recording, DiVA will stream back a response! This might take a second."
         )
         state = 2
     return (
         gr.Button(
-            value="Click to compare models!", interactive=True, variant="primary"
+            value="Click to run inference!", interactive=True, variant="primary"
         ),
         state,
     )
 
 
-def responses_complete(state):
-    if state == 2:
-        gr.Info(
-            "Give us your feedback! Mark which model gave you the best response so we can understand the quality of these different voice assistant models. NOTE: This will save an (irreversible) hash of your inputs to deduplicate any repeated votes."
-        )
-        state = 3
-    return state
-
-
 def clear_factory(button_id):
-    def clear(audio_input, text_prompt, model_order):
-        if button_id != None:
-            sr, y = audio_input
-            with scheduler.lock:
-                db.insert(
-                    {
-                        "audio_hash": hash(str(y)),
-                        "text_prompt": hash(text_prompt),
-                        "best": model_shorthand[model_order[button_id]],
-                    }
-                )
-        if anonymous:
-            random.shuffle(model_order)
+    def clear(audio_input, model_order):
         return (
             model_order,
             gr.Button(
                 value="Record Audio to Submit!",
                 interactive=False,
             ),
-            gr.Button(visible=False),
-            gr.Button(visible=False),
-            gr.Button(visible=False),
             None,
-            gr.Textbox(visible=False),
-            gr.Textbox(visible=False),
-            gr.Textbox(visible=False),
+            None
         )
 
     return clear
@@ -368,9 +156,8 @@ theme = gr.themes.Soft(
     neutral_hue="stone",
 )
 
-
-model_names = ["Llama 3 DiVA", "SALMONN", "Qwen Audio"]
-model_shorthand = ["via", "salmonn", "qwen"]
+model_names = ["DiVA Llama 3 8B"]
+model_shorthand = ["diva"]
 with gr.Blocks(theme=theme) as demo:
     state = gr.State(0)
     model_order = gr.State([0, 1])
@@ -378,26 +165,12 @@ with gr.Blocks(theme=theme) as demo:
         audio_input = gr.Audio(
             sources=["microphone"], streaming=False, label="Audio Input"
         )
-    with gr.Row():
-        prompt = gr.Textbox(
-            value="",
-            label="Text Prompt",
-            placeholder="Optional: Additional text prompt to influence how the model responds to your speech. e.g. 'Respond in a Haiku style.' or 'Translate the input to Arabic'",
-        )
 
     with gr.Row():
         btn = gr.Button(value="Record Audio to Submit!", interactive=False)
 
     with gr.Row():
-        with gr.Column(scale=1):
-            out1 = gr.Textbox(visible=False)
-            best1 = gr.Button(value="This response is best", visible=False)
-        with gr.Column(scale=1):
-            out2 = gr.Textbox(visible=False)
-            best2 = gr.Button(value="This response is best", visible=False)
-        with gr.Column(scale=1):
-            out3 = gr.Textbox(visible=False)
-            best3 = gr.Button(value="This response is best", visible=False)
+        out1 = gr.Textbox(visible=False)
 
     audio_input.stop_recording(
         recording_complete,
@@ -413,31 +186,16 @@ with gr.Blocks(theme=theme) as demo:
     )
     btn.click(
         fn=transcribe,
-        inputs=[audio_input, prompt, state, model_order],
-        outputs=[btn, out1, out2, out3, best1, best2, best3, state],
-    )
-    best1.click(
-        fn=clear_factory(0),
-        inputs=[audio_input, prompt, model_order],
-        outputs=[model_order, btn, best1, best2, best3, audio_input, out1, out2, out3],
-    )
-    best2.click(
-        fn=clear_factory(1),
-        inputs=[audio_input, prompt, model_order],
-        outputs=[model_order, btn, best1, best2, best3, audio_input, out1, out2, out3],
-    )
-    best3.click(
-        fn=clear_factory(2),
-        inputs=[audio_input, prompt, model_order],
-        outputs=[model_order, btn, best1, best2, best3, audio_input, out1, out2, out3],
+        inputs=[audio_input, state, model_order],
+        outputs=[btn, out1, state],
     )
     audio_input.clear(
         clear_factory(None),
-        [audio_input, prompt, model_order],
-        [model_order, btn, best1, best2, best3, audio_input, out1, out2, out3],
+        [audio_input, model_order],
+        [model_order, btn, audio_input, out1],
     )
     demo.load(
         fn=on_page_load, inputs=[state, model_order], outputs=[state, model_order]
     )
 
-demo.launch(share=True)
+demo.launch(share=True)