128 lines
3.9 KiB
Python
128 lines
3.9 KiB
Python
import sounddevice as sd
|
|
import numpy as np
|
|
import pyperclip
|
|
import requests
|
|
import sys
|
|
import argparse
|
|
from faster_whisper import WhisperModel
|
|
|
|
import os
|
|
os.environ["CT2_CUDA_ALLOW_FP16"] = "1"
|
|
|
|
# --- Configuration ---
|
|
MODEL_SIZE = "medium" # Options: "base", "small", "medium", "large-v3"
|
|
OLLAMA_URL = "http://localhost:11434/api/generate" # Default is 11434
|
|
DEFAULT_OLLAMA_MODEL = "qwen3:latest"
|
|
|
|
# Load Whisper on GPU
|
|
# float16 is faster and uses less VRAM on NVIDIA cards
|
|
print("Loading Whisper model...")
|
|
|
|
|
|
try:
|
|
model = WhisperModel(MODEL_SIZE, device="cuda", compute_type="float16")
|
|
except Exception as e:
|
|
print(f"Error loading GPU: {e}")
|
|
print("Falling back to CPU (Check your CUDA/cuDNN installation)")
|
|
model = WhisperModel(MODEL_SIZE, device="cuda", compute_type="int16")
|
|
|
|
|
|
def record_audio():
|
|
fs = 16000
|
|
print("\n[READY] Press Enter to START recording...")
|
|
input()
|
|
print("[RECORDING] Press Enter to STOP...")
|
|
|
|
recording = []
|
|
|
|
def callback(indata, frames, time, status):
|
|
if status:
|
|
print(status, file=sys.stderr)
|
|
recording.append(indata.copy())
|
|
|
|
with sd.InputStream(samplerate=fs, channels=1, callback=callback):
|
|
input()
|
|
|
|
return np.concatenate(recording, axis=0)
|
|
|
|
|
|
def main():
|
|
# 1. Setup Parser
|
|
print(f"System active. Model: {DEFAULT_OLLAMA_MODEL}")
|
|
parser = argparse.ArgumentParser(description="Whisper + Ollama CLI")
|
|
|
|
# Known Arguments (Hardcoded logic)
|
|
parser.add_argument("--nollm", "-n", action='store_true',
|
|
help="turn off llm")
|
|
parser.add_argument("--system", "-s", default=None,
|
|
help="The system prompt for Ollama")
|
|
parser.add_argument("--model_size", default="base",
|
|
help="Whisper model size: base, small, medium")
|
|
parser.add_argument(
|
|
"--ollama_model", default=DEFAULT_OLLAMA_MODEL, help="Ollama model name")
|
|
parser.add_argument(
|
|
"--num_ctx", default='5000', help="context length")
|
|
parser.add_argument(
|
|
"--temp", default='0.7', help="temperature")
|
|
|
|
# 2. Capture "Unknown" arguments
|
|
# args = known values, unknown = a list like ['--num_ctx', '4096', '--temp', '0.7']
|
|
args, unknown = parser.parse_known_args()
|
|
|
|
# Convert unknown list to a dictionary for the Ollama 'options' field
|
|
# This logic pairs ['--key', 'value'] into {key: value}
|
|
extra_options = {}
|
|
for i in range(0, len(unknown), 2):
|
|
key = unknown[i].lstrip('-') # remove the '--'
|
|
val = unknown[i+1]
|
|
# Try to convert numbers to actual ints/floats
|
|
try:
|
|
val = float(val) if '.' in val else int(val)
|
|
except ValueError:
|
|
pass
|
|
extra_options[key] = val
|
|
|
|
while True:
|
|
try:
|
|
audio_data = record_audio()
|
|
|
|
print("[TRANSCRIBING]...")
|
|
segments, _ = model.transcribe(audio_data.flatten(), beam_size=5)
|
|
text = "".join([segment.text for segment in segments]).strip()
|
|
|
|
if not text:
|
|
print("No speech detected. Try again.")
|
|
continue
|
|
|
|
# print(f"You said: {text}")
|
|
pyperclip.copy(text)
|
|
|
|
if (args.nollm == False):
|
|
# Send to Ollama
|
|
print(f"[OLLAMA] Thinking...")
|
|
payload = {
|
|
"model": args.ollama_model,
|
|
"prompt": text,
|
|
"stream": False,
|
|
"options": extra_options,
|
|
}
|
|
|
|
if args.system:
|
|
payload["system"] = args
|
|
response = requests.post(OLLAMA_URL, json=payload)
|
|
|
|
result = response.json().get("response", "")
|
|
print(f"\nLLM Response:\n{result}\n")
|
|
else:
|
|
print(f"\n{text}\n")
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nExiting...")
|
|
break
|
|
except Exception as e:
|
|
print(f"An error occurred: {e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|