Update app.py

app.py

@@ -3,12 +3,19 @@ import spaces
 import torch
 import subprocess
 import sys
+from threading import Thread
+from transformers import OlmoeForCausalLM, AutoTokenizer, TextIteratorStreamer
+import numpy as np
+import librosa
 
 # Install required packages
 subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "--force-reinstall", "--no-deps", "einops", "accelerate", "git+https://github.com/Muennighoff/transformers.git@olmoe"])
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-from transformers import OlmoeForCausalLM, AutoTokenizer
+# Import speech-to-speech components
+from VAD.vad_handler import VADHandler
+from STT.whisper_stt_handler import WhisperSTTHandler
+from TTS.parler_handler import ParlerTTSHandler
 
 model_name = "allenai/OLMoE-1B-7B-0924-Instruct"
 
@@ -18,12 +25,12 @@ try:
     model = OlmoeForCausalLM.from_pretrained(
         model_name, 
         trust_remote_code=True, 
-        torch_dtype=torch.float16,  # Using float16 for lower precision
+        torch_dtype=torch.float16,
         low_cpu_mem_usage=True,
         device_map="auto",
-        _attn_implementation="flash_attention_2"  # Enable Flash Attention 2
+        _attn_implementation="flash_attention_2"
     ).to(DEVICE)
-    model.gradient_checkpointing_enable()  # Enable gradient checkpointing
+    model.gradient_checkpointing_enable()
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 except Exception as e:
     print(f"Error loading model: {e}")
@@ -35,10 +42,35 @@ system_prompt = ("Adopt the persona of hilariously pissed off Andrej Karpathy "
                  "while always answering questions in full first principles analysis type of thinking "
                  "without using any analogies and always showing full working code or output in his answers.")
 
+# Setup speech-to-speech components
+from queue import Queue
+from threading import Event
+
+stop_event = Event()
+should_listen = Event()
+vad = VADHandler(stop_event, Queue(), Queue(), setup_args=(should_listen,))
+stt = WhisperSTTHandler(stop_event, Queue(), Queue())
+tts = ParlerTTSHandler(stop_event, Queue(), Queue(), setup_args=(should_listen,))
+
+@spaces.GPU
+def speech_to_text(audio):
+    if audio is None:
+        return ""
+    audio_np = librosa.resample(audio[1], orig_sr=audio[0], target_sr=16000)
+    audio_np = (audio_np * 32768).astype(np.int16)
+    
+    vad_output = vad.process(audio_np)
+    stt_output, _ = next(stt.process(vad_output))
+    return stt_output
+
+@spaces.GPU
+def user(user_message, history):
+    return "", history + [[user_message, None]]
+
 @spaces.GPU
-def generate_response(message, history, temperature, max_new_tokens):
+def bot(history, temperature, max_new_tokens):
     if model is None or tokenizer is None:
-        yield "Model or tokenizer not loaded properly. Please check the logs."
+        yield history
         return
 
     messages = [{"role": "system", "content": system_prompt}]
@@ -46,29 +78,36 @@ def generate_response(message, history, temperature, max_new_tokens):
         messages.append({"role": "user", "content": user_msg})
         if assistant_msg:
             messages.append({"role": "assistant", "content": assistant_msg})
-    messages.append({"role": "user", "content": message})
 
     inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
     
     try:
-        with torch.no_grad():
-            generated_ids = model.generate(
-                inputs,
-                max_new_tokens=max_new_tokens,
-                do_sample=True,
-                temperature=temperature,
-                eos_token_id=tokenizer.eos_token_id,
-            )
-        
-        generated_text = tokenizer.decode(generated_ids[0, inputs.shape[1]:], skip_special_tokens=True)
-        yield generated_text.strip()
-    except RuntimeError as e:
-        if "CUDA out of memory" in str(e):
-            yield "GPU memory exceeded. Try reducing the max tokens or using a smaller model."
-        else:
-            yield f"An error occurred: {str(e)}"
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = dict(
+            inputs=inputs,
+            streamer=streamer,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=temperature,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        generated_text = ""
+        for new_text in streamer:
+            generated_text += new_text
+            history[-1][1] = generated_text
+            yield history
+
     except Exception as e:
-        yield f"An unexpected error occurred: {str(e)}"
+        history[-1][1] = f"An error occurred: {str(e)}"
+        yield history
+
+def text_to_speech(text):
+    audio_output = np.concatenate(list(tts.process(text)))
+    return (16000, audio_output)
 
 css = """
   #output {
@@ -79,30 +118,30 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE (Now with Flash Attention 2!)")
+    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE (Now with Speech-to-Speech!)")
     chatbot = gr.Chatbot(elem_id="output")
+    audio_input = gr.Audio(source="microphone", type="numpy")
+    audio_output = gr.Audio()
     msg = gr.Textbox(label="Meow")
     with gr.Row():
         temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
         max_new_tokens = gr.Slider(minimum=50, maximum=4000, value=2000, step=50, label="Max New Tokens")
     clear = gr.Button("Clear")
 
-    def user(user_message, history):
-        return "", history + [[user_message, None]]
-
-    def bot(history, temp, max_tokens):
-        user_message = history[-1][0]
-        bot_message = ""
-        for token in generate_response(user_message, history[:-1], temp, max_tokens):
-            bot_message = token
-            history[-1][1] = bot_message
-            yield history
+    def process_audio(audio, history, temp, max_tokens):
+        text = speech_to_text(audio)
+        history = history + [[text, None]]
+        for new_history in bot(history, temp, max_tokens):
+            yield new_history, text_to_speech(new_history[-1][1])
 
+    audio_input.stop_recording(process_audio, [audio_input, chatbot, temperature, max_new_tokens], [chatbot, audio_output])
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
         bot, [chatbot, temperature, max_new_tokens], chatbot
+    ).then(
+        lambda history: text_to_speech(history[-1][1]), chatbot, audio_output
     )
     clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo.queue(api_open=True, max_size=10)  # Limiting queue size
-    demo.launch(debug=True, show_api=True, share=False)  # Disabled sharing for security
+    demo.queue(api_open=True, max_size=10)
+    demo.launch(debug=True, show_api=True, share=False)