ストリーミング

.gitattributes copy

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

ml_api/api.py

@@ -1,5 +1,9 @@
+from threading import Thread
+
 import torch
+from django.http import StreamingHttpResponse
 from ninja import NinjaAPI
+from transformers import TextIteratorStreamer
 
 from .model_loader import get_model
 from .schemas import ChatInput, ChatOutput
@@ -53,3 +57,59 @@ def chat(request, data: ChatInput):
 
     # ChatOutputスキーマに合わせてdictを返す
     return {"result": response_text}
+
+
+# ストリーミング用
+@api.post("/chat/stream")
+def chat_stream(request, data: ChatInput):
+    """
+    Qwenモデルを使用したストリーミングチャットAPI
+    """
+    user_input = data.text
+    model, tokenizer = get_model()
+
+    # 1. 会話フォーマットの作成
+    messages = [
+        {
+            "role": "system",
+            "content": "あなたは親切でフレンドリーなAIアシスタント「qwen」です。自然な日本語で簡潔に返事をしてください。",
+        },
+        {"role": "user", "content": user_input},
+    ]
+
+    # 2. プロンプトへの変換
+    text = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    # 3. ストリーマーの準備
+    # skip_prompt=True にしないと、質問文も一緒に返ってきてしまいます
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
+
+    # generateに渡す引数を準備
+    generation_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+    )
+
+    # 4. 別スレッドで生成を開始
+    # model.generateはブロッキング処理なので、スレッドに逃がさないと
+    # ストリーミング(yield)が開始されません。
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # 5. ジェネレーター関数の定義
+    def event_stream():
+        # streamerはイテレータとして動作し、新しいトークンが生成されるたびにループが回る
+        for new_text in streamer:
+            yield new_text
+
+    # StreamingHttpResponseにジェネレーターを渡して返す
+    return StreamingHttpResponse(event_stream(), content_type="text/event-stream")