from fastapi import FastAPI
from fastapi.responses import StreamingResponse, HTMLResponse
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import torch
import threading
app = FastAPI()
MODEL_NAME = "unsloth/gemma-3-1b-pt"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map="cpu",
torch_dtype=torch.float32,
low_cpu_mem_usage=True
)
class Request(BaseModel):
message: str
# =========================
# FRONTEND HTML (TON CODE)
# =========================
@app.get("/", response_class=HTMLResponse)
def home():
return """
Server Status
"""
# =========================
# CHAT STREAM (FIXÉ)
# =========================
@app.post("/chat-stream")
def chat_stream(req: Request):
# ✅ IMPORTANT : format chat propre
messages = [
{"role": "user", "content": req.message}
]
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt=True,
skip_special_tokens=True
)
def generate():
model.generate(
**inputs,
max_new_tokens=200,
temperature=0.4,
top_p=0.9,
repetition_penalty=1.2,
do_sample=True,
streamer=streamer
)
thread = threading.Thread(target=generate)
thread.start()
def token_stream():
for text in streamer:
yield text
return StreamingResponse(token_stream(), media_type="text/plain")