Update README.md

README.md

@@ -27,13 +27,34 @@ UniSS supports English and Chinese now.
 - **Demo:** https://cmots.github.io/uniss.github.io
 
 ## Quick Start
-1. Install the environment
+1. Install the environment and get the code
 ```bash
 conda create -n uniss python=3.10.16
 conda activate uniss
-pip install uniss
+git clone https://github.com/cmots/UniSS.git
+cd UniSS
+pip install -r requirements.txt
+# If you are in mainland China, you can set the mirror as follows:
+pip install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/ --trusted-host=mirrors.aliyun.com
 ```
-2. Run the code
+2. Download the weight
+The weight of UniSS is on [HuggingFace](https://huggingface.co/cmots/UniSS). 
+
+You have to download the model manually, you can download it via provided script:
+``` 
+python download_weight.py
+```
+
+or download via git clone (skip this if you have download via python script):
+``` bash
+mkdir -p pretrained_models
+
+# Make sure you have git-lfs installed (https://git-lfs.com)
+git lfs install
+
+git clone https://huggingface.co/cmots/UniSS pretrained_models/UniSS
+```
+3. Run the code
 ``` python
 import soundfile
 from uniss import UniSSTokenizer
@@ -41,44 +62,58 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from uniss import process_input, process_output
 
+# 1. Set the device, wav path, model path
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 wav_path = "prompt_audio.wav"
+model_path = "pretrained_models/UniSS"
 
-model_path = "cmots/UniSS"
+# 2. Set the mode and target language
+mode = 'Quality'    # 'Quality' or 'Performance'
+tgt_lang = "<|eng|>"    # for English output
+# tgt_lang = "<|cmn|>"  # for Chinese output
 
-# load the model, text tokenizer, and speech tokenizer
-model = AutoModelForCausalLM.from_pretrained(model_path)
+# 3. load the model, text tokenizer, and speech tokenizer
+model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-speech_tokenizer = UniSSTokenizer.from_pretrained(model_path)
 
-# extract speech tokens
+speech_tokenizer = UniSSTokenizer.from_pretrained(model_path, device=device)
+
+# 4. extract speech tokens
 glm4_tokens, bicodec_tokens = speech_tokenizer.tokenize(wav_path)
 
-tgt_lang = "<|eng|>"
 
-# process the input
-input_text = process_input(glm4_tokens, bicodec_tokens, "Quality", tgt_lang)
+# 5. process the input
+input_text = process_input(glm4_tokens, bicodec_tokens, mode, tgt_lang)
+input_token_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
 
-# translate the speech
+# 6. translate the speech
 output = model.generate(
-    glm4_tokens,
-    bicodec_tokens,
-    max_new_tokens=100,
-    num_beams=1,
-    early_stopping=True,
+    input_token_ids,
+    max_new_tokens=1500,
+    temperature=0.8,
+    top_p=0.8,
+    repetition_penalty=1.1
 )
-output_text = tokenizer.decode(output, skip_special_tokens=True)
 
-audio, translation, transcription = process_output(output_text, input_text, speech_tokenizer, "Quality", device)
+# 7. decode the output
+output_text = tokenizer.batch_decode(output, skip_special_tokens=True)
 
+# 8. process the output
+audio, translation, transcription = process_output(output_text[0], input_text, speech_tokenizer, mode, device)
+
+# 9. save and show the results
 soundfile.write("output_audio.wav", audio, 16000)
-print(translation)
-print(transcription)
 
+if mode == 'Quality':
+    print("Transcription:\n", transcription)
+print("Translation:\n", translation)
 ```
 
+More examples and details is on [Our Github Repo](https://github.com/cmots/UniSS).
+
 ## Citation
+If you find our paper and code useful in your research, please consider giving a like and citation.
 ```bibtex
 
 ```