Update space

LLM/mlx_language_model.py

@@ -0,0 +1,97 @@
+import logging
+from LLM.chat import Chat
+from baseHandler import BaseHandler
+from mlx_lm import load, stream_generate, generate
+from rich.console import Console
+import torch
+
+logging.basicConfig(
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+console = Console()
+
+
+class MLXLanguageModelHandler(BaseHandler):
+ """
+ Handles the language model part.
+ """
+
+ def setup(
+ self,
+ model_name="microsoft/Phi-3-mini-4k-instruct",
+ device="mps",
+ torch_dtype="float16",
+ gen_kwargs={},
+ user_role="user",
+ chat_size=1,
+ init_chat_role=None,
+ init_chat_prompt="You are a helpful AI assistant.",
+ ):
+ self.model_name = model_name
+ self.model, self.tokenizer = load(self.model_name)
+ self.gen_kwargs = gen_kwargs
+
+ self.chat = Chat(chat_size)
+ if init_chat_role:
+ if not init_chat_prompt:
+ raise ValueError(
+ "An initial promt needs to be specified when setting init_chat_role."
+ )
+ self.chat.init_chat({"role": init_chat_role, "content": init_chat_prompt})
+ self.user_role = user_role
+
+ self.warmup()
+
+ def warmup(self):
+ logger.info(f"Warming up {self.__class__.__name__}")
+
+ dummy_input_text = "Write me a poem about Machine Learning."
+ dummy_chat = [{"role": self.user_role, "content": dummy_input_text}]
+
+ n_steps = 2
+
+ for _ in range(n_steps):
+ prompt = self.tokenizer.apply_chat_template(dummy_chat, tokenize=False)
+ generate(
+ self.model,
+ self.tokenizer,
+ prompt=prompt,
+ max_tokens=self.gen_kwargs["max_new_tokens"],
+ verbose=False,
+ )
+
+ def process(self, prompt):
+ logger.debug("infering language model...")
+
+ self.chat.append({"role": self.user_role, "content": prompt})
+
+ # Remove system messages if using a Gemma model
+ if "gemma" in self.model_name.lower():
+ chat_messages = [
+ msg for msg in self.chat.to_list() if msg["role"] != "system"
+ ]
+ else:
+ chat_messages = self.chat.to_list()
+
+ prompt = self.tokenizer.apply_chat_template(
+ chat_messages, tokenize=False, add_generation_prompt=True
+ )
+ output = ""
+ curr_output = ""
+ for t in stream_generate(
+ self.model,
+ self.tokenizer,
+ prompt,
+ max_tokens=self.gen_kwargs["max_new_tokens"],
+ ):
+ output += t
+ curr_output += t
+ if curr_output.endswith((".", "?", "!", "<|end|>")):
+ yield curr_output.replace("<|end|>", "")
+ curr_output = ""
+ generated_text = output.replace("<|end|>", "")
+ torch.mps.empty_cache()
+
+ self.chat.append({"role": "assistant", "content": generated_text})

app.py

@@ -1,8 +1,43 @@
+# import base64
+# import pathlib
+# import tempfile
 import gradio as gr
+
+# recorder_js = pathlib.Path('recorder.js').read_text()
+# main_js = pathlib.Path('main.js').read_text()
+# record_button_js = pathlib.Path('record_button.js').read_text().replace('let recorder_js = null;', recorder_js).replace(
+# 'let main_js = null;', main_js)
+
+
+# def save_base64_video(base64_string):
+# base64_video = base64_string
+# video_data = base64.b64decode(base64_video)
+# with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
+# temp_filename = temp_file.name
+# temp_file.write(video_data)
+# print(f"Temporary MP4 file saved as: {temp_filename}")
+# return temp_filename
+# import os
+
+# os.system('python -m unidic download')
 from transformers import pipeline
 import numpy as np
 from VAD.vad_iterator import VADIterator
 import torch
+import librosa
+from mlx_lm import load, stream_generate, generate
+from LLM.chat import Chat
+from lightning_whisper_mlx import LightningWhisperMLX
+from melo.api import TTS
+
+LM_model, LM_tokenizer = load("mlx-community/SmolLM-360M-Instruct")
+chat = Chat(2)
+chat.init_chat({"role": "system", "content": "You are a helpful and friendly AI assistant. You are polite, respectful, and aim to provide concise responses of less than 20 words."})
+user_role = "user"
+
+tts_model = TTS(language="EN_NEWEST", device="auto")
+speaker_id = tts_model.hps.data.spk2id["EN-Newest"]
+blocksize = 512
 
 def int2float(sound):
  """
@@ -16,10 +51,13 @@ def int2float(sound):
  sound = sound.squeeze() # depends on the use case
  return sound
 
-min_speech_ms=500,
-max_speech_ms=float("inf"),
+text_str=""
+audio_output = None
+min_speech_ms=500
+max_speech_ms=float("inf")
+ASR_model = LightningWhisperMLX(model="distil-large-v3", batch_size=6, quant=None)
 transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base.en")
-vad_model, _ = torch.hub.load("snakers4/silero-vad", "silero_vad")
+vad_model, _ = torch.hub.load("snakers4/silero-vad:v4.0", "silero_vad")
 vad_iterator = VADIterator(
  vad_model,
  threshold=0.3,
@@ -31,131 +69,72 @@ vad_iterator = VADIterator(
 
 def transcribe(stream, new_chunk):
  sr, y = new_chunk
- print(sr)
- print(y.shape)
+ global text_str
+ global chat
+ global user_role
+ global audio_output
+ 
  audio_int16 = np.frombuffer(y, dtype=np.int16)
  audio_float32 = int2float(audio_int16)
+ audio_float32=librosa.resample(audio_float32, orig_sr=sr, target_sr=16000)
+ sr=16000
+ print(sr)
+ print(audio_float32.shape)
  vad_output = vad_iterator(torch.from_numpy(audio_float32))
+ 
  if vad_output is not None and len(vad_output) != 0:
- logger.debug("VAD: end of speech detected")
+ print("VAD: end of speech detected")
  array = torch.cat(vad_output).cpu().numpy()
  duration_ms = len(array) / sr * 1000
  if (not(duration_ms < min_speech_ms or duration_ms > max_speech_ms)):
- y = array.astype(np.float32)
- y /= np.max(np.abs(y))
- return stream, transcriber({"sampling_rate": sr, "raw": y})["text"] , y
- return stream, None , None
+ prompt=ASR_model.transcribe(array)["text"].strip()
+ chat.append({"role": user_role, "content": prompt})
+ chat_messages = chat.to_list()
+ prompt = LM_tokenizer.apply_chat_template(
+ chat_messages, tokenize=False, add_generation_prompt=True
+ )
+ output = generate(
+ LM_model,
+ LM_tokenizer,
+ prompt,
+ max_tokens=128,
+ )
+ # import pdb;pdb.set_trace()
+ generated_text = output.replace("<|end|>", "")
+ torch.mps.empty_cache()
+
+ chat.append({"role": "assistant", "content": generated_text})
+ text_str=generated_text
+ # import pdb;pdb.set_trace()
+ audio_chunk = tts_model.tts_to_file(text_str, speaker_id, quiet=True)
+ audio_chunk = (audio_chunk * 32768).astype(np.int16)
+ audio_output=(44100, audio_chunk)
+ # else:
+ # audio_output=None
+ text_str1=text_str
+ 
+ return stream, text_str1, audio_output
 
 demo = gr.Interface(
  transcribe,
  ["state", gr.Audio(sources=["microphone"], streaming=True, waveform_options=gr.WaveformOptions(sample_rate=16000))],
- ["state", "text", gr.Audio(label="Output", streaming=True, autoplay=True, waveform_options=gr.WaveformOptions(sample_rate=16000))],
+ ["state", "text", gr.Audio(label="Output", autoplay=True)],
  live=True,
 )
-
+# with demo:
+# start_button = gr.Button("Record Screen 🔴")
+# video_component = gr.Video(interactive=True, show_share_button=True, include_audio=True)
+
+
+# def toggle_button_label(returned_string):
+# if returned_string.startswith("Record"):
+# return gr.Button(value="Stop Recording ⚪"), None
+# else:
+# try:
+# temp_filename = save_base64_video(returned_string)
+# except Exception as e:
+# return gr.Button(value="Record Screen 🔴"), gr.Warning(f'Failed to convert video to mp4:\n{e}')
+# return gr.Button(value="Record Screen 🔴"), gr.Video(value=temp_filename, interactive=True,
+# show_share_button=True)
+# start_button.click(toggle_button_label, start_button, [start_button, video_component], js=record_button_js)
 demo.launch()
-# from transformers import pipeline
-# import torch
-
-# device = "cuda:0" if torch.cuda.is_available() else "cpu"
-
-# classifier = pipeline(
-# "audio-classification", model="MIT/ast-finetuned-speech-commands-v2", device=device
-# )
-
-# from transformers.pipelines.audio_utils import ffmpeg_microphone_live
-
-
-# def launch_fn(
-# wake_word="marvin",
-# prob_threshold=0.5,
-# chunk_length_s=2.0,
-# stream_chunk_s=0.25,
-# debug=False,
-# ):
-# if wake_word not in classifier.model.config.label2id.keys():
-# raise ValueError(
-# f"Wake word {wake_word} not in set of valid class labels, pick a wake word in the set {classifier.model.config.label2id.keys()}."
-# )
-
-# sampling_rate = classifier.feature_extractor.sampling_rate
-
-# mic = ffmpeg_microphone_live(
-# sampling_rate=sampling_rate,
-# chunk_length_s=chunk_length_s,
-# stream_chunk_s=stream_chunk_s,
-# )
-
-# print("Listening for wake word...")
-# for prediction in classifier(mic):
-# prediction = prediction[0]
-# if debug:
-# print(prediction)
-# if prediction["label"] == wake_word:
-# if prediction["score"] > prob_threshold:
-# return True
-
-# transcriber = pipeline(
-# "automatic-speech-recognition", model="openai/whisper-base.en", device=device
-# )
-# import sys
-
-
-# def transcribe(chunk_length_s=5.0, stream_chunk_s=1.0):
-# sampling_rate = transcriber.feature_extractor.sampling_rate
-
-# mic = ffmpeg_microphone_live(
-# sampling_rate=sampling_rate,
-# chunk_length_s=chunk_length_s,
-# stream_chunk_s=stream_chunk_s,
-# )
-
-# print("Start speaking...")
-# for item in transcriber(mic, generate_kwargs={"max_new_tokens": 128}):
-# sys.stdout.write("\033[K")
-# print(item["text"], end="\r")
-# if not item["partial"][0]:
-# break
-
-# return item["text"]
-
-# from huggingface_hub import HfFolder
-# import requests
-
-
-# def query(text, model_id="tiiuae/falcon-7b-instruct"):
-# api_url = f"https://api-inference.huggingface.co/models/{model_id}"
-# headers = {"Authorization": f"Bearer {HfFolder().get_token()}"}
-# payload = {"inputs": text}
-
-# print(f"Querying...: {text}")
-# response = requests.post(api_url, headers=headers, json=payload)
-# return response.json()[0]["generated_text"][len(text) + 1 :]
-
-# from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
-
-# processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-
-# model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts").to(device)
-# vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan").to(device)
-
-# from datasets import load_dataset
-
-# embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-# speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
-
-# def synthesise(text):
-# inputs = processor(text=text, return_tensors="pt")
-# speech = model.generate_speech(
-# inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder
-# )
-# return speech.cpu()
-
-
-# if __name__ == "__main__":
-# launch_fn(debug=True)
-# # transcription = transcribe()
-# # response = query(transcription)
-# # audio = synthesise(response)
- 
-# # Audio(audio, rate=16000, autoplay=True)

flagged/log.csv

@@ -0,0 +1,2 @@
+stream,new_chunk,stream,output 1,flag,username,timestamp
+,flagged/new_chunk/65327197a5439319f87d/audio.wav,,,,,2024-09-07 15:26:52.280189

main.js

@@ -0,0 +1,74 @@
+// main.js
+if (!ScreenCastRecorder.isSupportedBrowser()) {
+ console.error("Screen Recording not supported in this browser");
+}
+let recorder;
+let outputBlob;
+const stopRecording = () => __awaiter(void 0, void 0, void 0, function* () {
+ let currentState = "RECORDING";
+ // We should do nothing if the user try to stop recording when it is not started
+ if (currentState === "OFF" || recorder == null) {
+ return;
+ }
+ // if (currentState === "COUNTDOWN") {
+ // this.setState({
+ // currentState: "OFF",
+ // })
+ // }
+ if (currentState === "RECORDING") {
+ if (recorder.getState() === "inactive") {
+ // this.setState({
+ // currentState: "OFF",
+ // })
+ console.log("Inactive");
+ }
+ else {
+ outputBlob = yield recorder.stop();
+ console.log("Done recording");
+ // this.setState({
+ // outputBlob,
+ // currentState: "PREVIEW_FILE",
+ // })
+ window.currentState = "PREVIEW_FILE";
+ const videoSource = URL.createObjectURL(outputBlob);
+ window.videoSource = videoSource;
+ const fileName = "recording";
+ const link = document.createElement("a");
+ link.setAttribute("href", videoSource);
+ link.setAttribute("download", `${fileName}.webm`);
+ link.click();
+ }
+ }
+});
+const startRecording = () => __awaiter(void 0, void 0, void 0, function* () {
+ const recordAudio = true;
+ recorder = new ScreenCastRecorder({
+ recordAudio,
+ onErrorOrStop: () => stopRecording(),
+ });
+ try {
+ yield recorder.initialize();
+ }
+ catch (e) {
+ console.warn(`ScreenCastRecorder.initialize error: ${e}`);
+ // this.setState({ currentState: "UNSUPPORTED" })
+ window.currentState = "UNSUPPORTED";
+ return;
+ }
+ // this.setState({ currentState: "COUNTDOWN" })
+ const hasStarted = recorder.start();
+ if (hasStarted) {
+ // this.setState({
+ // currentState: "RECORDING",
+ // })
+ console.log("Started recording");
+ window.currentState = "RECORDING";
+ }
+ else {
+ stopRecording().catch(err => console.warn(`withScreencast.stopRecording threw an error: ${err}`));
+ }
+});
+
+// Set global functions to window.
+window.startRecording = startRecording;
+window.stopRecording = stopRecording;

mlx_models/distil-large-v3/config.json

@@ -0,0 +1,13 @@
+{
+ "n_mels": 128,
+ "n_audio_ctx": 1500,
+ "n_audio_state": 1280,
+ "n_audio_head": 20,
+ "n_audio_layer": 32,
+ "n_vocab": 51866,
+ "n_text_ctx": 448,
+ "n_text_state": 1280,
+ "n_text_head": 20,
+ "n_text_layer": 2,
+ "model_type": "whisper"
+}

mlx_models/distil-large-v3/weights.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fd01bf050289525f91ff3d96e2880381367a34beb3520ad516181517b209ebc
+size 1509130112

record_button.js

@@ -0,0 +1,40 @@
+// Setup if needed and start recording.
+async () => {
+ // Set up recording functions if not already initialized
+ if (!window.startRecording) {
+ let recorder_js = null;
+ let main_js = null;
+ }
+
+ // Function to fetch and convert video blob to base64 using async/await without explicit Promise
+ async function getVideoBlobAsBase64(objectURL) {
+ const response = await fetch(objectURL);
+ if (!response.ok) {
+ throw new Error('Failed to fetch video blob.');
+ }
+
+ const blob = await response.blob();
+
+ const reader = new FileReader();
+ reader.readAsDataURL(blob);
+
+ return new Promise((resolve, reject) => {
+ reader.onloadend = () => {
+ if (reader.result) {
+ resolve(reader.result.split(',')[1]); // Return the base64 string (without data URI prefix)
+ } else {
+ reject('Failed to convert blob to base64.');
+ }
+ };
+ });
+ }
+
+ if (window.currentState === "RECORDING") {
+ await window.stopRecording();
+ const base64String = await getVideoBlobAsBase64(window.videoSource);
+ return base64String;
+ } else {
+ window.startRecording();
+ return "Record";
+ }
+}

recorder.js

@@ -0,0 +1,112 @@
+// recorder.js
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+ function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+ return new (P || (P = Promise))(function (resolve, reject) {
+ function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+ function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+ function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+ step((generator = generator.apply(thisArg, _arguments || [])).next());
+ });
+};
+const BLOB_TYPE = "video/webm";
+class ScreenCastRecorder {
+ /** True if the current browser likely supports screencasts. */
+ static isSupportedBrowser() {
+ return (navigator.mediaDevices != null &&
+ navigator.mediaDevices.getUserMedia != null &&
+ navigator.mediaDevices.getDisplayMedia != null &&
+ MediaRecorder.isTypeSupported(BLOB_TYPE));
+ }
+ constructor({ recordAudio, onErrorOrStop }) {
+ this.recordAudio = recordAudio;
+ this.onErrorOrStopCallback = onErrorOrStop;
+ this.inputStream = null;
+ this.recordedChunks = [];
+ this.mediaRecorder = null;
+ }
+ /**
+ * This asynchronous method will initialize the screen recording object asking
+ * for permissions to the user which are needed to start recording.
+ */
+ initialize() {
+ return __awaiter(this, void 0, void 0, function* () {
+ const desktopStream = yield navigator.mediaDevices.getDisplayMedia({
+ video: true,
+ });
+ let tracks = desktopStream.getTracks();
+ if (this.recordAudio) {
+ const voiceStream = yield navigator.mediaDevices.getUserMedia({
+ video: false,
+ audio: true,
+ });
+ tracks = tracks.concat(voiceStream.getAudioTracks());
+ }
+ this.recordedChunks = [];
+ this.inputStream = new MediaStream(tracks);
+ this.mediaRecorder = new MediaRecorder(this.inputStream, {
+ mimeType: BLOB_TYPE,
+ });
+ this.mediaRecorder.ondataavailable = e => this.recordedChunks.push(e.data);
+ });
+ }
+ getState() {
+ if (this.mediaRecorder) {
+ return this.mediaRecorder.state;
+ }
+ return "inactive";
+ }
+ /**
+ * This method will start the screen recording if the user has granted permissions
+ * and the mediaRecorder has been initialized
+ *
+ * @returns {boolean}
+ */
+ start() {
+ if (!this.mediaRecorder) {
+ console.warn(`ScreenCastRecorder.start: mediaRecorder is null`);
+ return false;
+ }
+ const logRecorderError = (e) => {
+ console.warn(`mediaRecorder.start threw an error: ${e}`);
+ };
+ this.mediaRecorder.onerror = (e) => {
+ logRecorderError(e);
+ this.onErrorOrStopCallback();
+ };
+ this.mediaRecorder.onstop = () => this.onErrorOrStopCallback();
+ try {
+ this.mediaRecorder.start();
+ }
+ catch (e) {
+ logRecorderError(e);
+ return false;
+ }
+ return true;
+ }
+ /**
+ * This method will stop recording and then return the generated Blob
+ *
+ * @returns {(Promise|undefined)}
+ * A Promise which will return the generated Blob
+ * Undefined if the MediaRecorder could not initialize
+ */
+ stop() {
+ if (!this.mediaRecorder) {
+ return undefined;
+ }
+ let resolver;
+ const promise = new Promise(r => {
+ resolver = r;
+ });
+ this.mediaRecorder.onstop = () => resolver();
+ this.mediaRecorder.stop();
+ if (this.inputStream) {
+ this.inputStream.getTracks().forEach(s => s.stop());
+ this.inputStream = null;
+ }
+ return promise.then(() => this.buildOutputBlob());
+ }
+ buildOutputBlob() {
+ return new Blob(this.recordedChunks, { type: BLOB_TYPE });
+ }
+}

requirements.txt

@@ -1,7 +1,11 @@
-huggingface_hub==0.22.2
+huggingface_hub==0.23.2
 transformers[sentencepiece]
-sentencepiece
+sentencepiece==0.2.0
 datasets
-huggingface_hub
 torch==2.4.0
-torchaudio
+torchaudio
+librosa
+lightning-whisper-mlx>=0.0.10
+mlx-lm>=0.14.0
+melotts @ git+https://github.com/andimarafioti/MeloTTS.git#egg=MeloTTS # made a copy of MeloTTS to have compatible versions of transformers
+sounddevice==0.5.0