SFM2001/spititout

1import base642import io3import os4import re5import tempfile6from functools import lru_cache7from pathlib import Path8from typing import Literal9 10import numpy as np11import soundfile as sf12import torch13import uvicorn14from fastapi import FastAPI, HTTPException15from fastapi.middleware.cors import CORSMiddleware16from fastapi.responses import FileResponse17from fastapi.staticfiles import StaticFiles18from huggingface_hub import hf_hub_download19from pydantic import BaseModel20from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline21from openai import OpenAI22 23LLM_API = os.getenv("LLM_API", "").strip()24LLM_API_BASE_URL = os.getenv("LLM_API_BASE_URL", "https://api.deepseek.com").strip()25LLM_API_MODEL = os.getenv("LLM_API_MODEL", "deepseek-v4-flash").strip()26 27LLM_BACKEND = os.getenv("LLM_BACKEND", "llamacpp").lower()28TEXT_MODEL = os.getenv("TEXT_MODEL", "Qwen/Qwen3-4B-Instruct-2507")29GGUF_MODEL_REPO = os.getenv("GGUF_MODEL_REPO", "Qwen/Qwen3-1.7B-GGUF")30GGUF_MODEL_FILE = os.getenv("GGUF_MODEL_FILE", "Qwen3-1.7B-Q4_K_M.gguf")31LLAMA_CPP_N_CTX = int(os.getenv("LLAMA_CPP_N_CTX", "4096"))32LLAMA_CPP_N_THREADS = int(os.getenv("LLAMA_CPP_N_THREADS", str(max(1, os.cpu_count() or 1))))33ASR_MODEL = os.getenv("ASR_MODEL", "openai/whisper-tiny")34KOKORO_LANG_CODE = os.getenv("KOKORO_LANG_CODE", "z")35KOKORO_VOICE = os.getenv("KOKORO_VOICE", "zf_xiaobei")36MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "220"))37 38VENTING_SYSTEM_INSTRUCTION = """39你是一个非常懂人性、充满热情的“情绪嘴替”伙伴。40你的头号任务是：和用户站在一起，陪他们宣泄。41 42规则：431. 不要讲大道理，不要劝大度。用户在生气时，道理是没用的。442. 使用感性、强烈、发泄性的词汇。如果用户在骂某人或某事，你要义愤填膺，表达出“这也太离谱了”、“我也是服了”这种情绪。453. 你的目标是让用户感到“有人懂我，有人替我出气”。464. 语气像一个铁哥们或闺蜜，语气词可以多一点。475. 遵守安全底线：不宣扬仇恨犯罪，不进行人身威胁，不鼓励现实伤害。486. 响应长度要多样化，不要每次都回差不多长度。49"""50 51GUIDING_SYSTEM_INSTRUCTION = """52你现在是一个睿智、温和且具有同理心的心理导师。53用户刚才已经发泄过情绪了，现在他们同意听听你的建议或开导。54 55规则：561. 语气平和、坚定、宽容。572. 从客观角度分析问题，帮用户找到除了生气之外的解决方法，或者心理上的和解点。583. 肯定用户刚才发泄情绪的必要性，然后引导他们向前看。594. 每次回答不要太长，要循序渐进。605. 响应长度要根据用户状态变化。61"""62 63 64class Message(BaseModel):65    role: Literal["user", "model"]66    text: str67    timestamp: int68    audio: str | None = None69    aiAudio: str | None = None70 71 72class ChatRequest(BaseModel):73    history: list[Message]74    mode: Literal["VENTING", "GUIDING"]75    audioBase64: str | None = None76 77 78class SpeechRequest(BaseModel):79    text: str80 81 82app = FastAPI(title="SPITITOUT HF Space")83app.add_middleware(84    CORSMiddleware,85    allow_origins=["*"],86    allow_credentials=True,87    allow_methods=["*"],88    allow_headers=["*"],89)90 91 92def _device() -> str:93    return "cuda" if torch.cuda.is_available() else "cpu"94 95 96@lru_cache(maxsize=1)97def get_llm():98    tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL, trust_remote_code=True)99    dtype = torch.float16 if torch.cuda.is_available() else torch.float32100    model = AutoModelForCausalLM.from_pretrained(101        TEXT_MODEL,102        dtype=dtype,103        device_map="auto" if torch.cuda.is_available() else None,104        trust_remote_code=True,105    )106    if not torch.cuda.is_available():107        model.to("cpu")108    model.eval()109    return tokenizer, model110 111 112@lru_cache(maxsize=1)113def get_llamacpp_llm():114    try:115        from llama_cpp import Llama116    except Exception as exc:117        raise RuntimeError(118            "llama-cpp-python is not installed correctly. Check requirements.txt and Space build logs."119        ) from exc120 121    model_path = hf_hub_download(repo_id=GGUF_MODEL_REPO, filename=GGUF_MODEL_FILE)122    return Llama(123        model_path=model_path,124        n_ctx=LLAMA_CPP_N_CTX,125        n_threads=LLAMA_CPP_N_THREADS,126        n_gpu_layers=-1 if torch.cuda.is_available() else 0,127        verbose=False,128    )129 130 131@lru_cache(maxsize=1)132def get_api_client():133    if not LLM_API:134        raise RuntimeError("LLM_API is not set.")135    return OpenAI(136        api_key=LLM_API,137        base_url=LLM_API_BASE_URL,138    )139 140 141def generate_reply_api(messages: list[dict[str, str]]) -> str:142    client = get_api_client()143 144    # API 模式也限制历史和输出，避免慢、贵、重复145    api_messages = [msg.copy() for msg in messages]146 147    response = client.chat.completions.create(148        model=LLM_API_MODEL,149        messages=api_messages,150        max_tokens=min(MAX_NEW_TOKENS, 220),151        temperature=0.85,152        top_p=0.9,153        stream=False,154        extra_body={155            "thinking": {"type": "disabled"}156        },157    )158 159    text = response.choices[0].message.content or ""160    return remove_thinking_blocks(text) or "我听到了，你继续说。"161 162 163@lru_cache(maxsize=1)164def get_asr():165    device_id = 0 if torch.cuda.is_available() else -1166    dtype = torch.float16 if torch.cuda.is_available() else torch.float32167    return pipeline(168        "automatic-speech-recognition",169        model=ASR_MODEL,170        torch_dtype=dtype,171        device=device_id,172    )173 174 175@lru_cache(maxsize=1)176def get_tts():177    try:178        from kokoro import KPipeline179    except Exception as exc:180        raise RuntimeError(181            "Kokoro TTS is not installed correctly. Check requirements.txt and Space build logs."182        ) from exc183 184    return KPipeline(lang_code=KOKORO_LANG_CODE)185 186 187def transcribe_audio(audio_base64: str) -> str:188    audio_bytes = base64.b64decode(audio_base64)189    with tempfile.NamedTemporaryFile(suffix=".webm", delete=True) as audio_file:190        audio_file.write(audio_bytes)191        audio_file.flush()192        result = get_asr()(audio_file.name)193    return str(result.get("text", "")).strip()194 195 196# def build_chat_messages(request: ChatRequest, transcript: str | None) -> list[dict[str, str]]:197#     system = VENTING_SYSTEM_INSTRUCTION if request.mode == "VENTING" else GUIDING_SYSTEM_INSTRUCTION198#     messages = [{"role": "system", "content": system}]199 200#     for index, msg in enumerate(request.history[-12:]):201#         content = msg.text202#         if transcript and index == len(request.history[-12:]) - 1 and msg.role == "user":203#             content = transcript if content == "🎤 语音消息" else f"{content}\n\n语音补充：{transcript}"204#         messages.append({205#             "role": "assistant" if msg.role == "model" else "user",206#             "content": content,207#         })208 209#     return messages210 211 212def build_chat_messages(request: ChatRequest, transcript: str | None) -> list[dict[str, str]]:213    system = VENTING_SYSTEM_INSTRUCTION if request.mode == "VENTING" else GUIDING_SYSTEM_INSTRUCTION214 215    system += """216额外规则：2171. 不要复述上一轮回答。2182. 不要使用和上一轮相同的开头。2193. 用户只发短句时，只针对这句短句回应，不要把旧话题整段重复。2204. 每次最多 2 到 4 句话。221"""222 223    messages = [{"role": "system", "content": system}]224 225    recent_history = request.history[-4:]226 227    for index, msg in enumerate(recent_history):228        content = msg.text229        if transcript and index == len(recent_history) - 1 and msg.role == "user":230            content = transcript if content == "🎤 语音消息" else f"{content}\n\n语音补充：{transcript}"231 232        messages.append({233            "role": "assistant" if msg.role == "model" else "user",234            "content": content,235        })236 237    return messages238 239def messages_to_prompt(messages: list[dict[str, str]]) -> str:240    prompt = []241    for msg in messages:242        role = "assistant" if msg["role"] == "assistant" else msg["role"]243        prompt.append(f"<|im_start|>{role}\n{msg['content']}<|im_end|>")244    prompt.append("<|im_start|>assistant\n")245    return "\n".join(prompt)246 247 248def remove_thinking_blocks(text: str) -> str:249    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL | re.IGNORECASE)250    return text.strip()251 252 253def generate_reply(messages: list[dict[str, str]]) -> str:254    if LLM_API:255        return generate_reply_api(messages)256    if LLM_BACKEND == "llamacpp":257        return generate_reply_llamacpp(messages)258    return generate_reply_transformers(messages)259 260 261def generate_reply_llamacpp(messages: list[dict[str, str]]) -> str:262    llm = get_llamacpp_llm()263    no_think_messages = [msg.copy() for msg in messages]264    for msg in reversed(no_think_messages):265        if msg["role"] == "user":266            msg["content"] = f"{msg['content']}\n/no_think"267            break268    prompt = messages_to_prompt(no_think_messages)269    output = llm(270        prompt,271        max_tokens=MAX_NEW_TOKENS,272        temperature=0.7,273        top_p=0.8,274        repeat_penalty=1.12,275        stop=["<|im_end|>", "<|endoftext|>"],276    )277    text = output["choices"][0]["text"]278    return remove_thinking_blocks(text) or "我听到了，你继续说。"279 280 281def generate_reply_transformers(messages: list[dict[str, str]]) -> str:282    tokenizer, model = get_llm()283    try:284        prompt = tokenizer.apply_chat_template(285            messages,286            tokenize=False,287            add_generation_prompt=True,288            enable_thinking=False,289        )290    except TypeError:291        prompt = tokenizer.apply_chat_template(292            messages,293            tokenize=False,294            add_generation_prompt=True,295        )296 297    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)298    with torch.inference_mode():299        output_ids = model.generate(300            **inputs,301            max_new_tokens=MAX_NEW_TOKENS,302            do_sample=True,303            temperature=0.85,304            top_p=0.9,305            pad_token_id=tokenizer.eos_token_id,306        )307    generated_ids = output_ids[0][inputs.input_ids.shape[-1]:]308    text = tokenizer.decode(generated_ids, skip_special_tokens=True)309    return remove_thinking_blocks(text) or "我听到了，你继续说。"310 311 312def synthesize_speech(text: str) -> str | None:313    if not text.strip():314        return None315 316    pipeline_tts = get_tts()317    chunks = []318    for _, _, audio in pipeline_tts(text[:500], voice=KOKORO_VOICE, speed=1.05):319        chunks.append(np.asarray(audio, dtype=np.float32))320    if not chunks:321        return None322 323    audio = np.concatenate(chunks)324    wav_io = io.BytesIO()325    sf.write(wav_io, audio, 24000, format="WAV")326    return base64.b64encode(wav_io.getvalue()).decode("utf-8")327 328 329@app.get("/api/health")330def health():331    return {332        "ok": True,333        "runtime": "api" if LLM_API else "local",334        "llm_backend": "deepseek_api" if LLM_API else "llamacpp",335        "llm_api_base_url": LLM_API_BASE_URL if LLM_API else None,336        "llm_api_model": LLM_API_MODEL if LLM_API else None,337        "text_model": TEXT_MODEL,338        "gguf_model_repo": GGUF_MODEL_REPO,339        "gguf_model_file": GGUF_MODEL_FILE,340        "asr_model": ASR_MODEL,341        "kokoro_lang_code": KOKORO_LANG_CODE,342        "kokoro_voice": KOKORO_VOICE,343        "device": _device(),344    }345 346 347@app.post("/api/chat")348def chat(request: ChatRequest):349    try:350        transcript = transcribe_audio(request.audioBase64) if request.audioBase64 else None351        messages = build_chat_messages(request, transcript)352        return {"text": generate_reply(messages), "transcript": transcript}353    except Exception as exc:354        raise HTTPException(status_code=500, detail=str(exc)) from exc355 356 357@app.post("/api/speech")358def speech(request: SpeechRequest):359    try:360        return {"audio": synthesize_speech(request.text)}361    except Exception as exc:362        raise HTTPException(status_code=500, detail=str(exc)) from exc363 364 365dist_dir = Path(__file__).parent / "dist"366if dist_dir.exists():367    app.mount("/assets", StaticFiles(directory=dist_dir / "assets"), name="assets")368 369 370@app.get("/{path:path}")371def frontend(path: str):372    requested = dist_dir / path373    if requested.is_file():374        return FileResponse(requested)375    index = dist_dir / "index.html"376    if index.exists():377        return FileResponse(index)378    return {"message": "Run npm run build before serving the Space frontend."}379 380 381if __name__ == "__main__":382    port = int(os.getenv("PORT", "7860"))383    uvicorn.run(app, host="0.0.0.0", port=port)384