support for both streaming and non streaming

Dockerfile

@@ -8,4 +8,4 @@ RUN pip install --no-cache-dir --upgrade -r /requirements.txt
 
 RUN useradd -m -u 1000 user
 
-CMD ["fastapi", "run", "api.py", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "api:app", "--host", "0.0.0.0", "--port", "7860"]

api.py

@@ -1,8 +1,8 @@
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, HTMLResponse
 from fastapi import FastAPI, HTTPException
 import logging
 
-from llm_backend import chat_with_model
+from llm_backend import chat_with_model, stream_with_model
 from schema import ChatRequest
 
 """
@@ -26,6 +26,7 @@ def chat_stream(request: ChatRequest):
     kwargs = {
         "max_tokens": request.max_tokens,
         "temperature": request.temperature,
+        "stream": True,
         "top_p": request.top_p,
         "min_p": request.min_p,
         "typical_p": request.typical_p,
@@ -40,7 +41,33 @@ def chat_stream(request: ChatRequest):
         "mirostat_eta": request.mirostat_eta,
     }
     try:
-        token_generator = chat_with_model(request.chat_history, request.model, kwargs)
+        token_generator = stream_with_model(request.chat_history, request.model, kwargs)
         return StreamingResponse(token_generator, media_type="text/plain")
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/chat")
+def chat(request: ChatRequest):
+    kwargs = {
+        "max_tokens": request.max_tokens,
+        "temperature": request.temperature,
+        "stream": False,
+        "top_p": request.top_p,
+        "min_p": request.min_p,
+        "typical_p": request.typical_p,
+        "frequency_penalty": request.frequency_penalty,
+        "presence_penalty": request.presence_penalty,
+        "repeat_penalty": request.repeat_penalty,
+        "top_k": request.top_k,
+        "seed": request.seed,
+        "tfs_z": request.tfs_z,
+        "mirostat_mode": request.mirostat_mode,
+        "mirostat_tau": request.mirostat_tau,
+        "mirostat_eta": request.mirostat_eta,
+    }
+    try:
+        output = chat_with_model(request.chat_history, request.model, kwargs)
+        return HTMLResponse(output, media_type="text/plain")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

llm_backend.py

@@ -19,8 +19,7 @@ def get_llm(model_name):
 
 def format_chat(chat_history: list[Message]):
     """
-    Formats chat history and user input into a single string
-    suitable for the model.
+    Formats chat history and user input into a single string suitable for the model.
     """
     messages = []
     for msg in chat_history:
@@ -29,13 +28,16 @@ def format_chat(chat_history: list[Message]):
     return "\n".join(messages) + "\nAssistant:"
 
 
-def chat_with_model(chat_history, model, kwargs: dict):
+default_kwargs = dict(
+    max_tokens=2048,
+    top_k=1,
+)
+
+
+def stream_with_model(chat_history, model, kwargs: dict):
     prompt = format_chat(chat_history)
 
-    default_kwargs = dict(
-        max_tokens=2048,
-        top_k=1,
-    )
+    llm = get_llm(model)
 
     forced_kwargs = dict(
         stop=["\nUser:", "\nAssistant:", "</s>"],
@@ -43,8 +45,6 @@ def chat_with_model(chat_history, model, kwargs: dict):
         stream=True,
     )
 
-    llm = get_llm(model)
-
     input_kwargs = {**default_kwargs, **kwargs, **forced_kwargs}
     response = llm.__call__(prompt, **input_kwargs)
 
@@ -52,6 +52,23 @@ def chat_with_model(chat_history, model, kwargs: dict):
         yield token["choices"][0]["text"]
 
 
+def chat_with_model(chat_history, model, kwargs: dict):
+    prompt = format_chat(chat_history)
+
+    llm = get_llm(model)
+
+    forced_kwargs = dict(
+        stop=["\nUser:", "\nAssistant:", "</s>"],
+        echo=False,
+        stream=False,
+    )
+
+    input_kwargs = {**default_kwargs, **kwargs, **forced_kwargs}
+    response = llm.__call__(prompt, **input_kwargs)
+
+    return response["choices"][0]["text"]
+
+
 # %% example input
 # kwargs = dict(
 #     temperature=1,

schema.py

@@ -37,6 +37,7 @@ class Message(BaseModel):
 class ChatRequest(BaseModel):
     chat_history: List[Message]
     model: Literal["llama3.2", "falcon-mamba", "mistral-nemo"] = "llama3.2"
+    stream: bool = False
     max_tokens: Optional[int] = 65536
     temperature: float = 0.8
     top_p: float = 0.95