app.py

import io
import base64
import numpy as np
import soundfile as sf
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from kokoro import KPipeline
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
from pydantic import BaseModel
import uvicorn

# ── Model loading ──────────────────────────────────────────────────────────────
MODEL_NAME = "HuggingFaceTB/SmolLM2-360M-Instruct"

print("Loading LLM...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)
model.eval()
print("✅ LLM loaded")

print("Loading TTS...")
tts = KPipeline(lang_code="a")
print("✅ TTS loaded")

# ── FastAPI ────────────────────────────────────────────────────────────────────
app = FastAPI(title="AI Text + TTS API")

class QuestionRequest(BaseModel):
    question: str
    voice: str = "af_heart"
    max_new_tokens: int = 120


def run_llm(question: str, max_new_tokens: int = 120) -> str:
    messages = [
        {"role": "system", "content": "You are a helpful NLP chatbot. Reply simply and shortly."},
        {"role": "user", "content": question}
    ]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )
    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


def run_tts(text: str, voice: str = "af_heart") -> np.ndarray:
    parts = []
    for _, _, audio in tts(text, voice=voice):
        parts.append(audio)
    return np.concatenate(parts) if parts else np.array([])


# ── API routes ─────────────────────────────────────────────────────────────────
@app.get("/", response_class=HTMLResponse)
@app.get("/ui", response_class=HTMLResponse)
def serve_ui():
    return HTMLResponse(content=UI_HTML)


@app.get("/health")
def health():
    return {"status": "healthy", "llm": MODEL_NAME, "tts": "kokoro"}


@app.post("/api/text")
def generate_text(req: QuestionRequest):
    if not req.question.strip():
        raise HTTPException(400, "Question cannot be empty")
    try:
        answer = run_llm(req.question, req.max_new_tokens)
        return JSONResponse({"question": req.question, "answer": answer})
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/api/tts")
def generate_tts(req: QuestionRequest):
    if not req.question.strip():
        raise HTTPException(400, "Question cannot be empty")
    try:
        answer = run_llm(req.question, req.max_new_tokens)
        audio = run_tts(answer, req.voice)
        if audio.size == 0:
            raise HTTPException(500, "TTS produced no audio")
        buf = io.BytesIO()
        sf.write(buf, audio, 24000, format="WAV")
        buf.seek(0)
        return StreamingResponse(
            buf,
            media_type="audio/wav",
            headers={"X-Answer-Text": answer}
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/api/both")
def generate_both(req: QuestionRequest):
    if not req.question.strip():
        raise HTTPException(400, "Question cannot be empty")
    try:
        answer = run_llm(req.question, req.max_new_tokens)
        audio = run_tts(answer, req.voice)
        audio_b64 = ""
        if audio.size > 0:
            buf = io.BytesIO()
            sf.write(buf, audio, 24000, format="WAV")
            audio_b64 = base64.b64encode(buf.getvalue()).decode()
        return JSONResponse({
            "question": req.question,
            "answer": answer,
            "audio_base64": audio_b64,
            "sample_rate": 24000
        })
    except Exception as e:
        raise HTTPException(500, str(e))


# ── Embedded UI ────────────────────────────────────────────────────────────────
UI_HTML = """<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8"/>
  <meta name="viewport" content="width=device-width, initial-scale=1.0"/>
  <title>AI Text + TTS</title>
  <style>
    *{box-sizing:border-box;margin:0;padding:0}
    body{font-family:'Segoe UI',sans-serif;background:#0f1117;color:#e0e0e0;
         min-height:100vh;display:flex;flex-direction:column;align-items:center;padding:40px 16px}
    h1{font-size:1.8rem;font-weight:700;color:#fff;margin-bottom:6px}
    .sub{color:#888;font-size:.95rem;margin-bottom:32px;text-align:center}
    .card{background:#1a1d27;border:1px solid #2a2d3e;border-radius:14px;
          padding:28px;width:100%;max-width:720px;margin-bottom:20px}
    label{display:block;font-size:.85rem;color:#aaa;margin-bottom:8px;
          text-transform:uppercase;letter-spacing:.05em}
    textarea{width:100%;background:#0f1117;border:1px solid #2a2d3e;border-radius:8px;
             color:#e0e0e0;font-size:1rem;padding:12px 14px;resize:vertical;
             min-height:90px;outline:none;transition:border-color .2s}
    textarea:focus{border-color:#ff7043}
    .btn-row{display:flex;gap:12px;margin-top:16px;flex-wrap:wrap}
    button{flex:1;padding:12px 20px;border:none;border-radius:8px;font-size:.95rem;
           font-weight:600;cursor:pointer;transition:opacity .2s,transform .1s}
    button:active{transform:scale(.97)}
    button:disabled{opacity:.4;cursor:not-allowed}
    #btn-text{background:#3a86ff;color:#fff}
    #btn-tts{background:#ff7043;color:#fff}
    #btn-both{background:#8b5cf6;color:#fff}
    #btn-clear{background:#2a2d3e;color:#aaa;flex:0 0 auto;padding:12px 16px}
    .result{background:#1a1d27;border:1px solid #2a2d3e;border-radius:14px;
            padding:24px;width:100%;max-width:720px;margin-bottom:20px;display:none}
    .result.show{display:block}
    .rlabel{font-size:.8rem;color:#888;text-transform:uppercase;
            letter-spacing:.05em;margin-bottom:10px}
    .atext{font-size:1.05rem;line-height:1.6;color:#e0e0e0;white-space:pre-wrap}
    audio{width:100%;margin-top:14px;border-radius:8px;accent-color:#ff7043}
    .spin{display:inline-block;width:16px;height:16px;border:2px solid rgba(255,255,255,.3);
          border-top-color:#fff;border-radius:50%;animation:sp .7s linear infinite;
          vertical-align:middle;margin-right:6px}
    @keyframes sp{to{transform:rotate(360deg)}}
    .err{color:#ff5555;font-size:.9rem;margin-top:10px;display:none}
    .api{background:#1a1d27;border:1px solid #2a2d3e;border-radius:14px;
         padding:24px;width:100%;max-width:720px}
    .api h2{font-size:1rem;color:#aaa;margin-bottom:14px;
            text-transform:uppercase;letter-spacing:.05em}
    .ep{display:flex;align-items:center;gap:10px;padding:10px 12px;
        background:#0f1117;border-radius:8px;margin-bottom:8px;font-size:.88rem}
    .m{font-weight:700;font-size:.78rem;padding:3px 8px;border-radius:5px;flex-shrink:0}
    .post{background:#ff704322;color:#ff7043}
    .get{background:#3a86ff22;color:#3a86ff}
    .path{color:#e0e0e0;font-family:monospace}
    .desc{color:#666;font-size:.82rem;margin-left:auto}
  </style>
</head>
<body>
  <h1>🎙️ AI Text + Text-to-Speech</h1>
  <p class="sub">SmolLM2 generates answers · Kokoro converts them to speech</p>

  <div class="card">
    <label>Your Question</label>
    <textarea id="q" placeholder="e.g. Explain zero-shot classification in simple words."></textarea>
    <p class="err" id="err"></p>
    <div class="btn-row">
      <button id="btn-text"  onclick="query('text')">📝 Text Only</button>
      <button id="btn-tts"   onclick="query('tts')">🔊 Text + Audio</button>
      <button id="btn-both"  onclick="query('both')">⚡ Full JSON</button>
      <button id="btn-clear" onclick="clearAll()">✕ Clear</button>
    </div>
  </div>

  <div class="result" id="result">
    <div class="rlabel">Response</div>
    <div class="atext" id="ans"></div>
    <audio id="player" controls style="display:none"></audio>
  </div>

  <div class="api">
    <h2>API Endpoints</h2>
    <div class="ep"><span class="m get">GET</span><span class="path">/health</span><span class="desc">Model status</span></div>
    <div class="ep"><span class="m post">POST</span><span class="path">/api/text</span><span class="desc">JSON → text answer</span></div>
    <div class="ep"><span class="m post">POST</span><span class="path">/api/tts</span><span class="desc">JSON → WAV audio stream</span></div>
    <div class="ep"><span class="m post">POST</span><span class="path">/api/both</span><span class="desc">JSON → text + base64 audio</span></div>
  </div>

<script>
  const BTNS = ['btn-text','btn-tts','btn-both'];
  const LABELS = {text:'📝 Text Only', tts:'🔊 Text + Audio', both:'⚡ Full JSON'};

  function setLoading(mode){
    BTNS.forEach(id => document.getElementById(id).disabled = true);
    document.getElementById('btn-'+mode).innerHTML = '<span class="spin"></span>Generating…';
    document.getElementById('err').style.display = 'none';
  }

  function resetBtns(){
    BTNS.forEach(id => {
      const b = document.getElementById(id);
      b.disabled = false;
      b.innerHTML = LABELS[id.replace('btn-','')];
    });
  }

  function showErr(msg){
    const e = document.getElementById('err');
    e.textContent = msg; e.style.display = 'block';
  }

  function clearAll(){
    document.getElementById('q').value = '';
    document.getElementById('result').classList.remove('show');
    document.getElementById('ans').textContent = '';
    document.getElementById('player').style.display = 'none';
    document.getElementById('err').style.display = 'none';
    resetBtns();
  }

  function showResult(text, audioUrl){
    document.getElementById('ans').textContent = text;
    const p = document.getElementById('player');
    if(audioUrl){ p.src = audioUrl; p.style.display = 'block'; p.play(); }
    else { p.style.display = 'none'; }
    document.getElementById('result').classList.add('show');
  }

  async function query(mode){
    const question = document.getElementById('q').value.trim();
    if(!question){ showErr('Please enter a question first.'); return; }
    setLoading(mode);
    const body = JSON.stringify({question, voice:'af_heart'});
    const headers = {'Content-Type':'application/json'};
    try {
      if(mode === 'text'){
        const r = await fetch('/api/text',{method:'POST',headers,body});
        if(!r.ok) throw new Error(await r.text());
        const d = await r.json();
        showResult(d.answer, null);

      } else if(mode === 'tts'){
        const r = await fetch('/api/tts',{method:'POST',headers,body});
        if(!r.ok) throw new Error(await r.text());
        const answer = decodeURIComponent(r.headers.get('X-Answer-Text') || '(see audio)');
        const blob = await r.blob();
        showResult(answer, URL.createObjectURL(blob));

      } else {
        const r = await fetch('/api/both',{method:'POST',headers,body});
        if(!r.ok) throw new Error(await r.text());
        const d = await r.json();
        let url = null;
        if(d.audio_base64){
          const bytes = Uint8Array.from(atob(d.audio_base64), c=>c.charCodeAt(0));
          url = URL.createObjectURL(new Blob([bytes],{type:'audio/wav'}));
        }
        showResult(d.answer, url);
      }
    } catch(e){ showErr('Error: '+e.message); }
    finally { resetBtns(); }
  }
</script>
</body>
</html>"""

# ── Entry point ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)