|
|
import os |
|
|
from typing import List, Literal, Optional |
|
|
import requests |
|
|
|
|
|
from fastapi import FastAPI, HTTPException |
|
|
from fastapi.middleware.cors import CORSMiddleware |
|
|
from fastapi.staticfiles import StaticFiles |
|
|
from pydantic import BaseModel |
|
|
|
|
|
APP_TITLE = "HF Chat (Fathom-R1-14B via API)" |
|
|
APP_VERSION = "0.2.0" |
|
|
|
|
|
MODEL_ID = os.getenv("MODEL_ID", "FractalAIResearch/Fathom-R1-14B") |
|
|
STATIC_DIR = os.getenv("STATIC_DIR", "/app/static") |
|
|
ALLOWED_ORIGINS = os.getenv("ALLOWED_ORIGINS", "") |
|
|
HF_API_TOKEN = os.getenv("HF_API_TOKEN") |
|
|
|
|
|
app = FastAPI(title=APP_TITLE, version=APP_VERSION) |
|
|
|
|
|
if ALLOWED_ORIGINS: |
|
|
origins = [o.strip() for o in ALLOWED_ORIGINS.split(",") if o.strip()] |
|
|
app.add_middleware( |
|
|
CORSMiddleware, |
|
|
allow_origins=origins, |
|
|
allow_credentials=True, |
|
|
allow_methods=["*"], |
|
|
allow_headers=["*"], |
|
|
) |
|
|
|
|
|
class Message(BaseModel): |
|
|
role: Literal["system", "user", "assistant"] |
|
|
content: str |
|
|
|
|
|
class ChatRequest(BaseModel): |
|
|
messages: List[Message] |
|
|
max_new_tokens: int = 512 |
|
|
temperature: float = 0.7 |
|
|
top_p: float = 0.95 |
|
|
repetition_penalty: Optional[float] = 1.0 |
|
|
stop: Optional[List[str]] = None |
|
|
|
|
|
class ChatResponse(BaseModel): |
|
|
reply: str |
|
|
model: str |
|
|
|
|
|
def messages_to_prompt(messages: List[Message]) -> str: |
|
|
parts = [] |
|
|
for m in messages: |
|
|
if m.role == "system": |
|
|
parts.append(f"System: {m.content}") |
|
|
elif m.role == "user": |
|
|
parts.append(f"User: {m.content}") |
|
|
else: |
|
|
parts.append(f"Assistant: {m.content}") |
|
|
parts.append("Assistant:") |
|
|
return "\n".join(parts) |
|
|
|
|
|
@app.get("/api/health") |
|
|
def health(): |
|
|
return {"status": "ok", "model": MODEL_ID, "source": "huggingface-inference-api"} |
|
|
|
|
|
@app.post("/api/chat", response_model=ChatResponse) |
|
|
def chat(req: ChatRequest): |
|
|
if not HF_API_TOKEN: |
|
|
raise HTTPException(status_code=500, detail="HF_API_TOKEN not set") |
|
|
|
|
|
if not req.messages: |
|
|
raise HTTPException(status_code=400, detail="messages cannot be empty") |
|
|
|
|
|
prompt = messages_to_prompt(req.messages) |
|
|
|
|
|
headers = { |
|
|
"Authorization": f"Bearer {HF_API_TOKEN}" |
|
|
} |
|
|
|
|
|
payload = { |
|
|
"inputs": prompt, |
|
|
"parameters": { |
|
|
"max_new_tokens": req.max_new_tokens, |
|
|
"temperature": req.temperature, |
|
|
"top_p": req.top_p, |
|
|
"repetition_penalty": req.repetition_penalty, |
|
|
"return_full_text": True, |
|
|
} |
|
|
} |
|
|
|
|
|
response = requests.post( |
|
|
f"https://api-inference.huggingface.co/models/{MODEL_ID}", |
|
|
headers=headers, |
|
|
json=payload |
|
|
) |
|
|
|
|
|
if response.status_code != 200: |
|
|
raise HTTPException(status_code=response.status_code, detail=response.text) |
|
|
|
|
|
result = response.json() |
|
|
if isinstance(result, list) and result and "generated_text" in result[0]: |
|
|
full = result[0]["generated_text"] |
|
|
reply = full[len(prompt):].strip() if full.startswith(prompt) else full |
|
|
else: |
|
|
reply = str(result) |
|
|
|
|
|
if not reply: |
|
|
reply = "(No response generated.)" |
|
|
|
|
|
return ChatResponse(reply=reply, model=MODEL_ID) |
|
|
|
|
|
if os.path.isdir(STATIC_DIR): |
|
|
app.mount("/", StaticFiles(directory=STATIC_DIR, html=True), name="static") |
