├── .gitignore
├── LLM
├── __pycache__
│ ├── chat.cpython-310.pyc
│ └── mlx_lm.cpython-310.pyc
├── chat.py
└── mlx_lm.py
├── README.md
├── STT
├── __pycache__
│ └── lightning_whisper_mlx_handler.cpython-310.pyc
└── lightning_whisper_mlx_handler.py
├── TTS
├── __pycache__
│ └── melotts.cpython-310.pyc
└── melotts.py
├── baseHandler.py
├── listen_and_play.py
├── local_audio_streamer.py
├── logo.png
├── requirements.txt
├── s2s_pipeline.py
└── utils.py
/.gitignore:
--------------------------------------------------------------------------------
1 | __pycache__
2 | tmp
--------------------------------------------------------------------------------
/LLM/__pycache__/chat.cpython-310.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/shi3z/speech-to-speech-japanese/d3b394287d07684a155905084df708e8d53fb7c6/LLM/__pycache__/chat.cpython-310.pyc
--------------------------------------------------------------------------------
/LLM/__pycache__/mlx_lm.cpython-310.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/shi3z/speech-to-speech-japanese/d3b394287d07684a155905084df708e8d53fb7c6/LLM/__pycache__/mlx_lm.cpython-310.pyc
--------------------------------------------------------------------------------
/LLM/chat.py:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 | class Chat:
5 | """
6 | Handles the chat using to avoid OOM issues.
7 | """
8 |
9 | def __init__(self, size):
10 | self.size = size
11 | self.init_chat_message = None
12 | # maxlen is necessary pair, since a each new step we add an prompt and assitant answer
13 | self.buffer = []
14 |
15 | def append(self, item):
16 | self.buffer.append(item)
17 | if len(self.buffer) == 2 * (self.size + 1):
18 | self.buffer.pop(0)
19 | self.buffer.pop(0)
20 |
21 | def init_chat(self, init_chat_message):
22 | self.init_chat_message = init_chat_message
23 |
24 | def to_list(self):
25 | if self.init_chat_message:
26 | return [self.init_chat_message] + self.buffer
27 | else:
28 | return self.buffer
29 |
--------------------------------------------------------------------------------
/LLM/mlx_lm.py:
--------------------------------------------------------------------------------
1 | import logging
2 | from LLM.chat import Chat
3 | from baseHandler import BaseHandler
4 | from mlx_lm import load, stream_generate, generate
5 | from rich.console import Console
6 | import torch
7 | logging.basicConfig(
8 | format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
9 | )
10 | logger = logging.getLogger(__name__)
11 |
12 | console = Console()
13 |
14 | class MLXLanguageModelHandler(BaseHandler):
15 | """
16 | Handles the language model part.
17 | """
18 |
19 | def setup(
20 | self,
21 | model_name,
22 | device="mps",
23 | torch_dtype="float16",
24 | gen_kwargs={},
25 | user_role="user",
26 | chat_size=1,
27 | init_chat_role=None,
28 | #init_chat_prompt="You are a helpful AI assistant.",
29 | init_chat_prompt="あなたは日本語に堪能な友人です。英語を使ってはいけません。全て日本語で回答します. You must answer in Japanese",
30 | ):
31 | #model_name="mlx-community/Mistral-7B-Instruct-v0.3-4bit"
32 | #model_name="mlx-community/mlx-community/Phi-3-mini-4k-instruct-4bit"
33 | #model_name="mlx-community/Llama-3-Swallow-8B-Instruct-v0.1-4bit"
34 | #model_name="mlx-community/Llama-3-Swallow-8B-Instruct-v0.1-8bit"
35 | model_name="mlx-community/Llama-3-Swallow-8B-Instruct-v0.1-8bit"
36 | print(model_name)
37 | self.model_name = model_name
38 | model_id = model_name#'microsoft/Phi-3-mini-4k-instruct'
39 | self.model, self.tokenizer = load(model_id)
40 | self.gen_kwargs = gen_kwargs
41 |
42 | self.chat = Chat(chat_size)
43 | #init_chat_role=None
44 | init_chat_prompt="あなたは日本語に堪能な友人です。英語を使ってはいけません。全て日本語で回答します. You must answer in Japanese"
45 | if init_chat_role:
46 | if not init_chat_prompt:
47 | raise ValueError(
48 | "An initial promt needs to be specified when setting init_chat_role."
49 | )
50 | self.chat.init_chat({"role": init_chat_role, "content": init_chat_prompt})
51 | self.user_role = user_role
52 |
53 | self.warmup()
54 |
55 | def warmup(self):
56 | logger.info(f"Warming up {self.__class__.__name__}")
57 | return
58 |
59 | dummy_input_text = "Write me a poem about Machine Learning."
60 | dummy_chat = [{"role": self.user_role, "content": dummy_input_text}]
61 |
62 | n_steps = 2
63 |
64 | for _ in range(n_steps):
65 | prompt = self.tokenizer.apply_chat_template(dummy_chat, tokenize=False)
66 | generate(self.model, self.tokenizer, prompt=prompt, max_tokens=self.gen_kwargs["max_new_tokens"], verbose=False)
67 |
68 |
69 | def process(self, prompt):
70 | logger.debug("infering language model...")
71 |
72 |
73 | self.chat.append({"role": self.user_role, "content": f"{prompt}, "})
74 | prompt = self.tokenizer.apply_chat_template(self.chat.to_list(), tokenize=False, add_generation_prompt=True)
75 | output = ""
76 | curr_output = ""
77 | print(self.chat.to_list())
78 | for t in stream_generate(self.model, self.tokenizer, prompt, max_tokens=self.gen_kwargs["max_new_tokens"]):
79 | output += t
80 | curr_output += t
81 | if curr_output.endswith(('.', '?','？',',','。','!','<|end|>','<|eot_id|>')):
82 | yield curr_output.replace('<|end|>', '')
83 | curr_output = ""
84 | print(f"cur:{curr_output}")
85 | generated_text = output.replace('<|end|>', '')
86 | print(f"generated_text:{generated_text}")
87 | torch.mps.empty_cache()
88 |
89 | self.chat.append({"role": "assistant", "content": generated_text})
90 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |