Upload 12 files

Dockerfile

@@ -0,0 +1,35 @@
+# TalkToDoc production image.
+# Built to be interchangeable across hosts: Hugging Face Spaces, Render
+# (Docker deploy), Fly.io, Cloud Run, or any other platform that runs a
+# standard Docker image. Nothing here is tied to one specific host.
+# Configuration comes entirely from environment variables set by
+# whichever platform runs it (PORT, DATABASE_URL, ANTHROPIC_API_KEY,
+# SECRET_KEY, ENVIRONMENT).
+
+FROM python:3.11-slim
+
+# Whisper needs ffmpeg to read audio files, this installs it at the
+# system level since it can't come from pip.
+RUN apt-get update && apt-get install -y --no-install-recommends ffmpeg \
+    && rm -rf /var/lib/apt/lists/*
+
+# Run as a non-root user, standard practice and specifically expected by
+# Hugging Face Spaces' Docker SDK.
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /home/user/app
+
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . .
+
+ENV ENVIRONMENT=production
+
+# 7860 is Hugging Face Spaces' default port. Other platforms (like
+# Render) set $PORT themselves, and this falls back to that automatically.
+EXPOSE 7860
+
+CMD gunicorn app:app --bind 0.0.0.0:${PORT:-7860} --timeout 120 --workers 1

Procfile

@@ -0,0 +1 @@
+web: gunicorn app:app --bind 0.0.0.0:$PORT --timeout 120 --workers 1

README.md

@@ -1,10 +1,23 @@
 ---
 title: TalkToDoc
-emoji: 📊
-colorFrom: yellow
-colorTo: blue
+emoji: 🩺
+colorFrom: green
+colorTo: gray
 sdk: docker
-pinned: false
+app_port: 7860
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# TalkToDoc
+
+A multilingual communication tool that helps patients and healthcare
+providers understand each other across English, Yoruba, Hausa, Igbo, and
+Nigerian Pidgin. Built for rural Nigerian healthcare settings where
+doctor shortages and language barriers contribute to misdiagnosis and
+communication breakdowns.
+
+This Space runs the full application: speech-to-text, translation,
+symptom summarization, and text-to-speech, in the patient's own
+language.
+
+Requires ANTHROPIC_API_KEY, SECRET_KEY, and DATABASE_URL to be set as
+secrets in this Space's settings.

app.py

@@ -0,0 +1,186 @@
+"""
+TalkToDoc Flask application.
+Wires together speech-to-text, translation, natural language understanding,
+and text-to-speech into the patient-provider communication flow described
+in the research document.
+
+Two-step flow:
+1. Patient submits input (text or audio) in their chosen language.
+   The system transcribes it if needed, translates it to English, and
+   summarizes the likely symptoms or intent for the provider.
+2. Provider submits a reply in English. The system translates it back
+   into the patient's language and generates spoken audio for it.
+"""
+
+import os
+import uuid
+from datetime import datetime, timezone
+
+from dotenv import load_dotenv
+from flask import Flask, request, jsonify, render_template, session
+
+import database
+import stt
+import translation
+import nlu
+import tts
+
+load_dotenv()
+
+AUDIO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static", "audio")
+
+
+def create_app():
+    flask_app = Flask(__name__)
+    flask_app.secret_key = os.environ["SECRET_KEY"]
+    os.makedirs(AUDIO_DIR, exist_ok=True)
+    return flask_app
+
+
+app = create_app()
+
+# Runs whether the app is started directly (python app.py, local development)
+# or imported by a production server like gunicorn (Render deployment).
+# Initialization code inside "if __name__ == '__main__'" only runs on a
+# direct start, gunicorn never triggers it, so anything the app needs to
+# work at all has to happen here instead.
+database.init_db()
+tts.preload_models()
+
+
+def get_current_user_id(language):
+    """
+    Returns the user_id for the current browser session, creating a new
+    user and session record the first time this browser is seen. This is
+    what keeps a patient's messages tied to the same person instead of
+    creating a brand new anonymous patient on every single submission.
+    """
+    if "user_id" not in session:
+        user_id = database.add_user(name="Patient", preferred_language=language, role="patient")
+        session_id = database.add_session(
+            user_id=user_id,
+            start_time=datetime.now(timezone.utc).isoformat(),
+        )
+        session["user_id"] = user_id
+        session["session_id"] = session_id
+
+    return session["user_id"]
+
+
+@app.route("/")
+def index():
+    return render_template("patient.html", languages=translation.SUPPORTED_LANGUAGES)
+
+
+@app.route("/provider")
+def provider_page():
+    return render_template("provider.html")
+
+
+@app.route("/patient-input", methods=["POST"])
+def patient_input():
+    language = request.form.get("language")
+    if not language or language.lower() not in translation.SUPPORTED_LANGUAGES:
+        return jsonify({"error": f"language must be one of {translation.SUPPORTED_LANGUAGES}"}), 400
+    language = language.lower()
+
+    text = request.form.get("text")
+    audio_file = request.files.get("audio")
+
+    if not text and not audio_file:
+        return jsonify({"error": "provide either text or audio"}), 400
+
+    if audio_file:
+        temp_path = os.path.join(AUDIO_DIR, f"upload_{uuid.uuid4().hex}.wav")
+        audio_file.save(temp_path)
+        try:
+            patient_text = stt.transcribe_audio(temp_path, language)
+        finally:
+            os.remove(temp_path)
+    else:
+        patient_text = text
+
+    english_text = translation.translate(patient_text, language, "english")
+    nlu_summary = nlu.interpret_query(english_text)
+
+    user_id = get_current_user_id(language)
+    interaction_id = database.add_interaction(
+        user_id=user_id,
+        input_text=patient_text,
+        detected_language=language,
+        translated_text=english_text,
+        nlu_summary=nlu_summary,
+        timestamp=datetime.now(timezone.utc).isoformat(),
+    )
+
+    return jsonify({
+        "interaction_id": interaction_id,
+        "patient_text": patient_text,
+        "translated_text": english_text,
+        "nlu_summary": nlu_summary,
+    })
+
+
+@app.route("/provider-response", methods=["POST"])
+def provider_response():
+    interaction_id = request.form.get("interaction_id")
+    response_text = request.form.get("response_text")
+
+    if not interaction_id or not response_text:
+        return jsonify({"error": "interaction_id and response_text are required"}), 400
+
+    interaction = database.get_interaction(interaction_id)
+    if not interaction:
+        return jsonify({"error": "interaction not found"}), 404
+
+    patient_language = interaction["detected_language"]
+    translated_response = translation.translate(response_text, "english", patient_language)
+
+    audio_filename = f"response_{interaction_id}.wav"
+    audio_path = os.path.join(AUDIO_DIR, audio_filename)
+    tts.synthesize_speech(translated_response, patient_language, audio_path)
+
+    database.update_interaction_response(interaction_id, response_text, translated_response)
+
+    return jsonify({
+        "translated_response": translated_response,
+        "audio_url": f"/static/audio/{audio_filename}",
+    })
+
+
+@app.route("/interaction/<interaction_id>")
+def get_interaction(interaction_id):
+    interaction = database.get_interaction(interaction_id)
+    if not interaction:
+        return jsonify({"error": "interaction not found"}), 404
+
+    audio_filename = f"response_{interaction_id}.wav"
+    audio_path = os.path.join(AUDIO_DIR, audio_filename)
+    interaction["audio_url"] = f"/static/audio/{audio_filename}" if os.path.exists(audio_path) else None
+
+    return jsonify(interaction)
+
+
+@app.route("/pending-interactions")
+def pending_interactions():
+    return jsonify(database.get_pending_interactions())
+
+
+@app.route("/history")
+def history():
+    if "user_id" not in session:
+        return jsonify([])
+    return jsonify(database.get_interactions_for_user(session["user_id"]))
+
+
+@app.route("/end-session", methods=["POST"])
+def end_session_route():
+    if "session_id" in session:
+        database.end_session(session["session_id"], datetime.now(timezone.utc).isoformat())
+    session.clear()
+    return jsonify({"status": "ended"})
+
+
+if __name__ == "__main__":
+    is_production = os.environ.get("ENVIRONMENT") == "production"
+    app.run(debug=not is_production, host="0.0.0.0", port=int(os.environ.get("PORT", 5000)))

database.py

@@ -0,0 +1,138 @@
+"""
+Database layer for TalkToDoc.
+Local development and testing use SQLite (no setup needed). When deployed,
+DATABASE_URL is set by Render and the app uses Postgres instead, since
+Render's free tier wipes local files like a SQLite database on every
+restart.
+
+Queries are written once using SQLite-style '?' placeholders and adapted
+automatically for Postgres, so there's a single query per function rather
+than two versions of everything.
+"""
+
+import os
+from pathlib import Path
+
+import sqlite3
+
+DB_PATH = Path(__file__).parent / "talktodoc.db"
+SCHEMA_PATH = Path(__file__).parent / "schema.sql"
+SCHEMA_PATH_POSTGRES = Path(__file__).parent / "schema_postgres.sql"
+
+DATABASE_URL = os.environ.get("DATABASE_URL")
+
+if DATABASE_URL:
+    import psycopg2
+    import psycopg2.extras
+
+
+def get_connection():
+    if DATABASE_URL:
+        return psycopg2.connect(DATABASE_URL, cursor_factory=psycopg2.extras.RealDictCursor)
+    connection = sqlite3.connect(DB_PATH)
+    connection.row_factory = sqlite3.Row
+    connection.execute("PRAGMA foreign_keys = ON")
+    return connection
+
+
+def _run(connection, query, params=()):
+    if DATABASE_URL:
+        query = query.replace("?", "%s")
+    cursor = connection.cursor()
+    cursor.execute(query, params)
+    return cursor
+
+
+def _insert_and_get_id(connection, query, params):
+    if DATABASE_URL:
+        cursor = _run(connection, query + " RETURNING id", params)
+        return cursor.fetchone()["id"]
+    cursor = _run(connection, query, params)
+    return cursor.lastrowid
+
+
+def init_db():
+    schema_path = SCHEMA_PATH_POSTGRES if DATABASE_URL else SCHEMA_PATH
+    schema_sql = schema_path.read_text()
+    with get_connection() as connection:
+        if DATABASE_URL:
+            connection.cursor().execute(schema_sql)
+        else:
+            connection.executescript(schema_sql)
+
+
+def add_user(name, preferred_language, role):
+    with get_connection() as connection:
+        return _insert_and_get_id(
+            connection,
+            "INSERT INTO app_user (name, preferred_language, role) VALUES (?, ?, ?)",
+            (name, preferred_language, role),
+        )
+
+
+def add_session(user_id, start_time):
+    with get_connection() as connection:
+        return _insert_and_get_id(
+            connection,
+            "INSERT INTO session (user_id, start_time) VALUES (?, ?)",
+            (user_id, start_time),
+        )
+
+
+def end_session(session_id, end_time):
+    with get_connection() as connection:
+        _run(connection, "UPDATE session SET end_time = ? WHERE id = ?", (end_time, session_id))
+
+
+def add_interaction(user_id, input_text, detected_language, translated_text, nlu_summary, timestamp):
+    with get_connection() as connection:
+        return _insert_and_get_id(
+            connection,
+            """INSERT INTO interaction
+               (user_id, input_text, detected_language, translated_text, nlu_summary, timestamp)
+               VALUES (?, ?, ?, ?, ?, ?)""",
+            (user_id, input_text, detected_language, translated_text, nlu_summary, timestamp),
+        )
+
+
+def get_interactions_for_user(user_id):
+    with get_connection() as connection:
+        cursor = _run(
+            connection,
+            "SELECT * FROM interaction WHERE user_id = ? ORDER BY timestamp",
+            (user_id,),
+        )
+        return [dict(row) for row in cursor.fetchall()]
+
+
+def get_pending_interactions():
+    with get_connection() as connection:
+        cursor = _run(
+            connection,
+            "SELECT * FROM interaction WHERE provider_response IS NULL ORDER BY timestamp",
+        )
+        return [dict(row) for row in cursor.fetchall()]
+
+
+def get_interaction(interaction_id):
+    with get_connection() as connection:
+        cursor = _run(connection, "SELECT * FROM interaction WHERE id = ?", (interaction_id,))
+        row = cursor.fetchone()
+        return dict(row) if row else None
+
+
+def update_interaction_response(interaction_id, provider_response, translated_response):
+    with get_connection() as connection:
+        _run(
+            connection,
+            "UPDATE interaction SET provider_response = ?, translated_response = ? WHERE id = ?",
+            (provider_response, translated_response, interaction_id),
+        )
+
+
+if __name__ == "__main__":
+    init_db()
+    if DATABASE_URL:
+        print("Database initialized (Postgres)")
+    else:
+        print("Database created at:", DB_PATH)

nlu.py

@@ -0,0 +1,71 @@
+"""
+Natural language understanding for TalkToDoc.
+Interprets basic health-related queries to identify symptoms, intent, or
+information requests, per functional requirement 5 in the research
+document.
+
+This is a communication aid, not a diagnostic tool. The document is
+explicit that the system does not provide formal medical diagnosis, so
+this module only summarizes what the patient is communicating, it never
+suggests a diagnosis or treatment.
+
+Uses the Claude API (Anthropic), same as translation.py.
+
+MOCK_MODE: if set to "true" in .env, this uses simple keyword matching
+instead of a real API call, so the rest of the app can be tested for
+free, with no API key. Not a substitute for testing real NLU quality.
+"""
+
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+MOCK_MODE = os.environ.get("MOCK_MODE", "").lower() == "true"
+
+if not MOCK_MODE:
+    from anthropic import Anthropic
+
+    client = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
+    MODEL = "claude-haiku-4-5-20251001"
+
+_MOCK_KEYWORDS = ["headache", "fever", "stomach", "cough", "dizzy", "pain", "vomit", "rash"]
+
+
+def interpret_query(english_text):
+    """
+    english_text: the patient's message, already translated to English
+    Returns a short plain-language summary of the likely symptoms, intent,
+    or information request, to help the provider quickly understand what
+    the patient needs. Not a diagnosis.
+    """
+    if MOCK_MODE:
+        text_lower = english_text.lower()
+        found = [word for word in _MOCK_KEYWORDS if word in text_lower]
+        symptoms = ", ".join(found) if found else "an unspecified concern"
+        return f"Patient reports {symptoms}. Requesting guidance. (Mock summary, no AI used.)"
+
+    prompt = (
+        "A patient sent the following message to a healthcare provider. "
+        "In two sentences or less, summarize the likely symptoms, intent, "
+        "or information request. Do not diagnose or suggest treatment, "
+        "only summarize what the patient is communicating.\n\n"
+        f"Message: {english_text}"
+    )
+
+    response = client.messages.create(
+        model=MODEL,
+        max_tokens=200,
+        messages=[{"role": "user", "content": prompt}],
+    )
+
+    return response.content[0].text.strip()
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 2:
+        print('Usage: python nlu.py "english text"')
+    else:
+        print(interpret_query(sys.argv[1]))

requirements.txt

@@ -0,0 +1,11 @@
+flask==3.1.3
+python-dotenv==1.2.2
+anthropic==0.117.0
+openai-whisper==20250625
+transformers==5.14.1
+torch==2.11.0
+torchaudio==2.11.0
+scipy==1.18.0
+yarngpt==0.2.0
+psycopg2-binary==2.9.12
+gunicorn==26.0.0

schema.sql

@@ -0,0 +1,27 @@
+CREATE TABLE IF NOT EXISTS app_user (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    name TEXT NOT NULL,
+    preferred_language TEXT NOT NULL,
+    role TEXT NOT NULL CHECK (role IN ('patient', 'provider'))
+);
+
+CREATE TABLE IF NOT EXISTS session (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    user_id INTEGER NOT NULL,
+    start_time TEXT NOT NULL,
+    end_time TEXT,
+    FOREIGN KEY (user_id) REFERENCES app_user (id)
+);
+
+CREATE TABLE IF NOT EXISTS interaction (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    user_id INTEGER NOT NULL,
+    input_text TEXT,
+    detected_language TEXT,
+    translated_text TEXT,
+    nlu_summary TEXT,
+    provider_response TEXT,
+    translated_response TEXT,
+    timestamp TEXT NOT NULL,
+    FOREIGN KEY (user_id) REFERENCES app_user (id)
+);

schema_postgres.sql

@@ -0,0 +1,25 @@
+CREATE TABLE IF NOT EXISTS app_user (
+    id SERIAL PRIMARY KEY,
+    name TEXT NOT NULL,
+    preferred_language TEXT NOT NULL,
+    role TEXT NOT NULL CHECK (role IN ('patient', 'provider'))
+);
+
+CREATE TABLE IF NOT EXISTS session (
+    id SERIAL PRIMARY KEY,
+    user_id INTEGER NOT NULL REFERENCES app_user (id),
+    start_time TEXT NOT NULL,
+    end_time TEXT
+);
+
+CREATE TABLE IF NOT EXISTS interaction (
+    id SERIAL PRIMARY KEY,
+    user_id INTEGER NOT NULL REFERENCES app_user (id),
+    input_text TEXT,
+    detected_language TEXT,
+    translated_text TEXT,
+    nlu_summary TEXT,
+    provider_response TEXT,
+    translated_response TEXT,
+    timestamp TEXT NOT NULL
+);

stt.py

@@ -0,0 +1,51 @@
+"""
+Speech-to-text for TalkToDoc.
+Converts a patient's spoken audio into text using a local Whisper model.
+Language is passed in from the manual language selector, not auto-detected.
+"""
+
+import whisper
+
+_model = None
+
+# Maps the app's language selection to Whisper's language codes.
+# Whisper has no dedicated Nigerian Pidgin code, so Pidgin is passed as
+# English, since Pidgin is close enough to English that this gives Whisper
+# the best chance of an accurate transcription.
+LANGUAGE_MAP = {
+    "english": "en",
+    "yoruba": "yo",
+    "hausa": "ha",
+    "igbo": "ig",
+    "pidgin": "en",
+}
+
+
+def _get_model():
+    global _model
+    if _model is None:
+        _model = whisper.load_model("base")
+    return _model
+
+
+def transcribe_audio(audio_path, language=None):
+    """
+    audio_path: path to an audio file (wav, mp3, m4a, etc.)
+    language: one of "english", "yoruba", "hausa", "igbo", "pidgin"
+    Returns the transcribed text.
+    """
+    model = _get_model()
+    whisper_language = LANGUAGE_MAP.get(language.lower()) if language else None
+    result = model.transcribe(audio_path, language=whisper_language)
+    return result["text"].strip()
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 2:
+        print("Usage: python stt.py <audio_file> [language]")
+    else:
+        audio_file = sys.argv[1]
+        selected_language = sys.argv[2] if len(sys.argv) > 2 else None
+        print(transcribe_audio(audio_file, selected_language))

translation.py

@@ -0,0 +1,69 @@
+"""
+Translation for TalkToDoc.
+Translates patient input into English for the provider, and translates the
+provider's reply back into the patient's selected language.
+
+Uses the Claude API (Anthropic), as specified in the research document's
+software requirements (3.5.4).
+
+MOCK_MODE: if set to "true" in .env, this skips the real API call
+entirely and returns the input text unchanged instead. This exists so
+the whole app (Whisper, YarnGPT, MMS-TTS, the database, all the routing
+and session logic) can be run and tested locally with zero cost and no
+API key, before deciding to actually acquire one. It is not a substitute
+for testing real translation quality, only for testing everything else.
+"""
+
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+MOCK_MODE = os.environ.get("MOCK_MODE", "").lower() == "true"
+
+if not MOCK_MODE:
+    from anthropic import Anthropic
+
+    client = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
+    MODEL = "claude-haiku-4-5-20251001"
+
+SUPPORTED_LANGUAGES = ["english", "yoruba", "hausa", "igbo", "pidgin"]
+
+
+def translate(text, source_language, target_language):
+    """
+    text: the text to translate
+    source_language: one of "english", "yoruba", "hausa", "igbo", "pidgin"
+    target_language: one of "english", "yoruba", "hausa", "igbo", "pidgin"
+    Returns the translated text only.
+    """
+    if MOCK_MODE:
+        if source_language.lower() == target_language.lower():
+            return text
+        return f"[MOCK, no real translation: {source_language} -> {target_language}] {text}"
+
+    prompt = (
+        f"Translate the following text from {source_language} to {target_language}. "
+        f"Reply with only the translated text and nothing else, no explanation.\n\n"
+        f"Text: {text}"
+    )
+
+    response = client.messages.create(
+        model=MODEL,
+        max_tokens=500,
+        messages=[{"role": "user", "content": prompt}],
+    )
+
+    return response.content[0].text.strip()
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 4:
+        print('Usage: python translation.py "text" source_language target_language')
+    else:
+        input_text = sys.argv[1]
+        source_lang = sys.argv[2]
+        target_lang = sys.argv[3]
+        print(translate(input_text, source_lang, target_lang))

tts.py

@@ -0,0 +1,138 @@
+"""
+Text-to-speech for TalkToDoc.
+Converts translated text responses into audible speech in the patient's
+selected language.
+
+English, Yoruba, Igbo, and Hausa use YarnGPT, a local model trained
+specifically on Nigerian-accented speech.
+
+Nigerian Pidgin uses MMS-TTS (facebook/mms-tts-pcm) instead, since YarnGPT
+has no Pidgin support. MMS-TTS is the only free option found with a
+dedicated Pidgin checkpoint.
+
+Note: importing this file triggers a one-time download of YarnGPT's audio
+tokenizer files. This is expected and only happens once.
+
+Performance note: YarnGPT's own generate_speech() function reloads its
+full model from disk on every call, which is too slow for a live app.
+This file loads the model once and reuses it, using the same generation
+steps YarnGPT's own function uses internally, just without the reload.
+Call preload_models() once when the app starts, so the first real request
+isn't slow either.
+"""
+
+import torch
+import torchaudio
+import scipy.io.wavfile
+from transformers import VitsModel, AutoTokenizer
+from yarngpt.core import load_model_and_tokenizer, SPEAKER_MAPPING, AVAILABLE_SPEAKERS
+
+# Default speaker used for each YarnGPT-supported language.
+YARNGPT_SPEAKERS = {
+    "english": "idera",
+    "yoruba": "abayomi",
+    "igbo": "chioma",
+    "hausa": "amina",
+}
+
+_yarngpt_model = None
+_yarngpt_tokenizer = None
+
+_pidgin_model = None
+_pidgin_tokenizer = None
+
+
+def _get_yarngpt_model():
+    global _yarngpt_model, _yarngpt_tokenizer
+    if _yarngpt_model is None:
+        _yarngpt_model, _yarngpt_tokenizer = load_model_and_tokenizer()
+    return _yarngpt_model, _yarngpt_tokenizer
+
+
+def _get_pidgin_model():
+    global _pidgin_model, _pidgin_tokenizer
+    if _pidgin_model is None:
+        _pidgin_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-pcm")
+        _pidgin_model = VitsModel.from_pretrained("facebook/mms-tts-pcm")
+    return _pidgin_model, _pidgin_tokenizer
+
+
+def preload_models():
+    """Loads both TTS backends into memory ahead of time. Call this once
+    when the Flask app starts, so the first real request isn't slow."""
+    _get_yarngpt_model()
+    _get_pidgin_model()
+
+
+def _generate_yarngpt_speech(text, speaker, language, temperature=0.1, repetition_penalty=1.1, max_length=4000):
+    """
+    Same steps as yarngpt's own generate_speech(), but reuses the model
+    already loaded by _get_yarngpt_model() instead of reloading it.
+    """
+    model_speaker = SPEAKER_MAPPING.get(speaker, speaker)
+    if model_speaker not in AVAILABLE_SPEAKERS:
+        raise ValueError(f"Unknown speaker: {speaker}")
+
+    model, audio_tokenizer = _get_yarngpt_model()
+
+    prompt = audio_tokenizer.create_prompt(text, language, model_speaker)
+    input_ids = audio_tokenizer.tokenize_prompt(prompt)
+    attention_mask = torch.ones_like(input_ids)
+
+    output = model.generate(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        do_sample=True,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_length=max_length,
+        pad_token_id=model.config.eos_token_id,
+        eos_token_id=model.config.eos_token_id,
+    )
+
+    codes = audio_tokenizer.get_codes(output)
+    audio = audio_tokenizer.get_audio(codes)
+    return audio
+
+
+def synthesize_speech(text, language, output_path):
+    """
+    text: the text to speak
+    language: one of "english", "yoruba", "hausa", "igbo", "pidgin"
+    output_path: where to save the resulting .wav file
+    Returns output_path.
+    """
+    language = language.lower()
+
+    if language in YARNGPT_SPEAKERS:
+        speaker = YARNGPT_SPEAKERS[language]
+        audio = _generate_yarngpt_speech(text, speaker=speaker, language=language)
+        torchaudio.save(output_path, audio, sample_rate=24000)
+        return output_path
+
+    if language == "pidgin":
+        model, tokenizer = _get_pidgin_model()
+        inputs = tokenizer(text, return_tensors="pt")
+        with torch.no_grad():
+            output = model(**inputs).waveform
+        scipy.io.wavfile.write(
+            output_path,
+            rate=model.config.sampling_rate,
+            data=output.numpy().squeeze(),
+        )
+        return output_path
+
+    raise ValueError(f"Unsupported language: {language}")
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 3:
+        print('Usage: python tts.py "text" language [output_file]')
+    else:
+        input_text = sys.argv[1]
+        selected_language = sys.argv[2]
+        output_file = sys.argv[3] if len(sys.argv) > 3 else "output.wav"
+        synthesize_speech(input_text, selected_language, output_file)
+        print("Saved to", output_file)