Update app.py

app.py

@@ -4,23 +4,13 @@ import torch
 import subprocess
 import sys
 from threading import Thread
-from transformers import AutoTokenizer, TextIteratorStreamer
-import numpy as np
-
-from queue import Queue
-from threading import Event
+from transformers import TextIteratorStreamer
 
 # Install required packages
-subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "--force-reinstall", "librosa", "parler_tts", "melotts", "funasr", "sounddevice", "nltk", "sounddevice", "deepfilternet", "einops", "accelerate", "ChatTTS", "git+https://github.com/Muennighoff/transformers.git@olmoe"])
+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)
 
-# Import the custom OlmoeForCausalLM
-from transformers import OlmoeForCausalLM
-import librosa
-# Import speech-to-speech components
-from VAD.vad_handler import VADHandler
-from STT.whisper_stt_handler import WhisperSTTHandler
-from TTS.parler_handler import ParlerTTSHandler
+from transformers import OlmoeForCausalLM, AutoTokenizer
 
 model_name = "allenai/OLMoE-1B-7B-0924-Instruct"
 
@@ -30,12 +20,12 @@ try:
     model = OlmoeForCausalLM.from_pretrained(
         model_name, 
         trust_remote_code=True, 
-        torch_dtype=torch.float16,
+        torch_dtype=torch.float16,  # Using float16 for lower precision
         low_cpu_mem_usage=True,
         device_map="auto",
-        _attn_implementation="flash_attention_2"
+        _attn_implementation="flash_attention_2"  # Enable Flash Attention 2
     ).to(DEVICE)
-    model.gradient_checkpointing_enable()
+    model.gradient_checkpointing_enable()  # Enable gradient checkpointing
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 except Exception as e:
     print(f"Error loading model: {e}")
@@ -47,32 +37,10 @@ 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
-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 bot(history, temperature, max_new_tokens):
+def generate_response(message, history, temperature, max_new_tokens):
     if model is None or tokenizer is None:
-        yield history
+        yield "Model or tokenizer not loaded properly. Please check the logs."
         return
 
     messages = [{"role": "system", "content": system_prompt}]
@@ -80,6 +48,7 @@ def bot(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)
     
@@ -97,53 +66,52 @@ def bot(history, temperature, max_new_tokens):
         thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
 
-        generated_text = ""
+        partial_message = ""
         for new_text in streamer:
-            generated_text += new_text
-            history[-1][1] = generated_text
-            yield history
-
+            partial_message += new_text
+            yield partial_message.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)}"
     except Exception as 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)
+        yield f"An unexpected error occurred: {str(e)}"
 
 css = """
   #output {
-    height: 1000px; 
+    height: 1100px; 
     overflow: auto; 
-    border: 2px solid #ccc; 
+    border: 3px solid #ccc; 
   }
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE (Now with Speech-to-Speech!)")
+    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE (Now with Flash Attention 2 and Streaming!)")
     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 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])
+    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
 
-    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)
-    demo.launch(debug=True, show_api=True, share=False)
+    demo.queue(api_open=True, max_size=10)  # Limiting queue size
+    demo.launch(debug=True, show_api=True, share=False)  # Disabled sharing for security