README update

README.md

@@ -27,14 +27,117 @@ tags:
 - Whisper
 ---
 
-# Causal Whisper
+# CarelessWhisper - Causal Whisper Streaming Model
+Causal Whisper Streaming is a fine tuned version of OpenAI Whisper, which can handle causal data and perform real-time transcription. 
 
-causal-whisper is an openai-whisper based streaming model.
-Available on 5 differenet chunk sizes for the **english transcription**:
-1. 1000[msec] (small, large-v2).
-1. 300[msec] (base, small, large-v2).
-1. 200[msec] (base, small, large-v2).
-1. 100[msec] (base, small, large-v2).
-1. 40[msec] (base, small, large-v2).
+![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)
+[![Demo on Hugging Face](https://img.shields.io/badge/🤗%20Demo-Hugging%20Face-blueviolet?logo=huggingface&logoColor=white)](https://huggingface.co/spaces/MLSpeech/CarelessWhisper-causal-streaming)
+
+
+## Setup
+We used Python 3.9.16, PyTorch 2.6.0, and PyTorch-Lightning 2.5.0 to train and test our models.
+All of the required dependencies are available on `requirements.txt` file. Make sure all of the packages are installed before running anything.
+Portions of this code are adapted from [OpenAI's Whisper](https://github.com/openai/whisper).
+
+## Available Models
+We fine-tuned three different sizes of Whisper, all support english only transcription.
+A `large-v2` that was fine tuned on multilingual data is available, and supports English, French, Spanish, German and Portuguese with chunk size of 300 miliseconds.
+
+| Size | Chunk Size [msec] | Multilingual | 
+|:----:|:-----------------:|:------------:|
+| base | 40, 100, 200, 300 |  N/A         |
+| small| 40, 100, 200, 300, 1000| N/A     |
+|large-v2| 40, 100, 200, 300, 1000| 300   |
+
+
+## Running Inference
+To run inference, download the repo content, and run accroding to following sections. 
+### CLI Usage
+The transcription model is easily activated using the next command:
+```bash
+# Using a local microphone for streaming transcription, dumping the recording to out.wav
+python causal_whisper_stream/streaming_transcribe.py \
+--output_filename out.wav \  # path to dump recording to - optional
+--channels 2 \  # number of channels
+--model small \   # model size
+--chunk_size 300 \  # chunk size in milliseconds
+--device cuda \  # inference device
+--beam_size 5 \  # run inference with beam of size 5 
+--ca_kv_cache \  # run inference with CA kv cache
+```
+
+A simulation of a stream on a wav file is also available:
+```bash
+# Simulating a stream on a wav file
+python causal_whisper_stream/streaming_transcribe.py \
+--model small \   # model size
+--chunk_size 300 \  # chunk size in milliseconds
+--device cuda \  # inference device
+--beam_size 5 \  # run inference with beam of size 5 
+--ca_kv_cache \  # run inference with CA kv cache
+--wav_file /path/to/audio.wav \  # path to audio file to transcribe
+--simulate_stream \  # a must for a simulation to occur
+--use_latency \  # simulating real latency of the chunk size - optional.
+```
+
+### Python Usage
+If you prefer using python, a code sinppet utilizing a microphone or a wav file is provided below:
+
+```python
+import torch
+import causal_whisper_stream
+
+model_size = "small" # model size
+chunk_size = 300 # chunk size in milliseconds
+multilingual = False # currently on large-v2_300msec supports other languages than english.
+device = "cuda" if torch.cuda.is_avaialable() else "cpu"
+
+model = causal_whisper_stream.load_streaming_model(name=model_size,
+                                                   gran=chunk_size,
+                                                   multilingual=multilingual,
+                                                   device=device)
+
+# using a local microphone recording 
+texts_microphone = model.transcribe(output_filename="/path/to/dump/file.wav",
+                         channels=2,
+                         beam_size=5,
+                         ca_kv_cache=True)
+
+# Simulating on a wav file
+texts_wav_simulation = model.transcribe(simulate_stream=True,
+                                        wav_file="/path/to/file/you/want/to/transcribe.wav",
+                                        beam_size=5,
+                                        ca_kv_cache=True)
+```
+
+## Training
+In order to train using LoRA, you can use our existing code. Make sure all the requirements are installed. 
+
+### CLI Interface
+```bash
+# 
+python training_code/train.py \
+--lora \  # use LoRA fine-tuning, currently only this way is supported.
+--streaming_train \  # streaming traing
+--simulate_stream \  # simulate a stream of data, use streaming mel
+--dataset LIBRI-960-ALIGNED \  # dataset name from a dict 
+--name example_training_base_model \  # name for differentiation
+--size base \  # model size to train on 
+--batch_size 32 \  # batch size
+--epochs 10 \  # maximal epochs
+--learning_rate 1e-5 \  # initial learning rate
+--rank 32 \  # LoRA rank 
+--gran 15 \  # Chunk size in encoder frames units. 15 encoder frames = 300msec
+--extra_gran_blocks 1 \  # extra blocks for inital chunk. 1 means that an extra block is added for the beginning. i.e., the inital chunk will be 600msec.
+--streaming_fraction 0.25 \  # fraction of samples out of the possible subsequences in the streaming process to train on
+--top_k 5 \  # how many ckpts to save.
+```
+
+For more options and training configurations, run:
+```bash
+python training_code/train.py --help
+```
+### 🙏 Acknowledgements
+
+This project uses components from [OpenAI's Whisper](https://github.com/openai/whisper), licensed under the MIT License.
 
-A large-v2 model is available for french, portuguese, spanish and german with a chunk size of 300[msec].