Update README.md

README.md

@@ -32,38 +32,53 @@ We continue to expand [Meta-Llama-3-8B-Instruct](https://huggingface.co/meta-lla
 
 ## How to Get Started with the Model
 
+First, we need to convert the audio file to sound tokens
+
 ```python
 import torch
 import torchaudio
 from encodec import EncodecModel
 from encodec.utils import convert_audio
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
 
-# Audio to Sound Tokens
 def audio_to_sound_tokens(audio_path, target_bandwidth=1.5, device="cuda"):
+    # Initialize Encodec
     model = EncodecModel.encodec_model_24khz()
     model.set_target_bandwidth(target_bandwidth)
     model.to(device)
 
+    # Load and preprocess audio
     wav, sr = torchaudio.load(audio_path)
     wav = convert_audio(wav, sr, model.sample_rate, model.channels)
     wav = wav.unsqueeze(0).to(device)
 
+    # Encode audio
     with torch.no_grad():
         encoded_frames = model.encode(wav)
     codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1)
 
+    # Flatten codes
     audio_code1, audio_code2 = codes[0][0], codes[0][1]
     flatten_tokens = torch.stack((audio_code1, audio_code2), dim=1).flatten().tolist()
 
+    # Convert to sound tokens
     result = ''.join(f'<|sound_{num}|>' for num in flatten_tokens)
     return f'<|sound_start|>{result}<|sound_end|>'
 
-# LLM Pipeline Setup
+# Usage
+sound_tokens = audio_to_sound_tokens("/path/to/your/audio/file")
+```
+
+Then, we can inference the model the same as any other LLM.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
+
 def setup_pipeline(model_path, use_4bit=True):
     tokenizer = AutoTokenizer.from_pretrained(model_path)
 
     model_kwargs = {"device_map": "auto"}
+    
     if use_4bit:
         model_kwargs["quantization_config"] = BitsAndBytesConfig(
             load_in_4bit=True,
@@ -73,9 +88,9 @@ def setup_pipeline(model_path, use_4bit=True):
         )
 
     model = AutoModelForCausalLM.from_pretrained(model_path, **model_kwargs)
+
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Text Generation
 def generate_text(pipe, messages, max_new_tokens=64, temperature=0.0, do_sample=False):
     generation_args = {
         "max_new_tokens": max_new_tokens,
@@ -83,26 +98,18 @@ def generate_text(pipe, messages, max_new_tokens=64, temperature=0.0, do_sample=
         "temperature": temperature,
         "do_sample": do_sample,
     }
+
     output = pipe(messages, **generation_args)
     return output[0]['generated_text']
 
-# Main process
-def audio_to_text(audio_path, model_path, use_4bit=True):
-    # Convert audio to sound tokens
-    sound_tokens = audio_to_sound_tokens(audio_path)
-    
-    # Setup LLM pipeline
-    pipe = setup_pipeline(model_path, use_4bit)
-    
-    # Generate text
-    messages = [{"role": "user", "content": sound_tokens}]
-    return generate_text(pipe, messages)
-
-# Usage example
-audio_path = "/path/to/your/audio/file"
-model_path = "jan-hq/Jan-Llama3-0708"
-
-generated_text = audio_to_text(audio_path, model_path)
+# Usage
+llm_path = "jan-hq/Jan-Llama3-0708"
+pipe = setup_pipeline(llm_path, use_4bit=True)
+messages = [
+    {"role": "user", "content": sound_tokens},
+]
+generated_text = generate_text(pipe, messages)
+print(generated_text)
 ```
 
 ## Training process