update app, add model param

Author: peng3307165

Dockerfile

@@ -3,21 +3,22 @@ FROM python:3.10
 ADD . /workspace/code-repo
 WORKDIR /workspace/code-repo
 
-RUN pip install fastapi uvicorn
+RUN pip3 install fastapi uvicorn
 RUN pip3 install requests
+RUN pip3 install jinja2
 
 ENV PYTHONPATH /workspace/code-repo
 
 RUN apt-get update && apt-get install -y curl
 
-RUN curl -fsSL https://ollama.com/install.sh | sh
-
-RUN ollama serve
-RUN ollama run llama3.2
+#RUN curl -fsSL https://ollama.com/install.sh | sh
+#RUN ollama serve
+#RUN ollama run llama3.2
 
 ENV FLASK_RUN_HOST=0.0.0.0
 
 EXPOSE 8000
 
 CMD sh -c "ollama serve & ollama run llama3.2"
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000", "--load-balancer", "sunrpc"]
+#CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000", "--load-balancer", "sunrpc"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

app.py

@@ -73,13 +73,25 @@ async def generate(request: Request):
     request_dict = await request.json()
     model = request_dict.get("model", "")
     messages = request_dict.get("messages", dict())
-    stream = request_dict.get("stream", False)
     prompt = messages[-1].get("content", "")
+    stream = request_dict.get("stream", False)
+
+    infer_param = request_dict.get("infer_param", dict())
+    max_tokens = infer_param.get('max_tokens', 4096)
+    temperature = infer_param.get('temperature', 0.9)
+    top_p = infer_param.get('top_p', 0.9)
+    n = infer_param.get('n', 1)
+    stop = infer_param.get('stop', [])
 
     data = {
         "model": model,
         "prompt": prompt,
-        "stream": stream
+        "stream": stream,
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "n": n,
+        "stop": stop,
     }
     # url_generate = "http://127.0.0.1:11434/api/generate"
 

example/ollama_app.py

@@ -8,7 +8,8 @@ def run():
     prompt = 'hello'
     model = 'llama3.2'
     messages = [{"role": "user", "content": prompt}]
-    data = {'model': model, 'messages': messages}
+    infer_param = {}
+    data = {'engine': 'ollama', 'model': model, 'messages': messages, 'infer_param':  infer_param}
     headers = {"Content-Type": "application/json"}
     response = requests.post(url, headers=headers, data=json.dumps(data))