migrated to groq -- suuuuuper fast!

.gitignore

@@ -1,5 +1,6 @@
 __pycache__
 venv
 .vscode
+.env
 # script for some housekeeping
-f.py            
+f.py

Dockerfile

@@ -1,9 +1,26 @@
-FROM ollama/ollama
+# FROM ollama/ollama
 
-RUN mkdir -p /.ollama && chmod 777 /.ollama
+# RUN mkdir -p /.ollama && chmod 777 /.ollama
 
-ENV OLLAMA_MAX_LOADED_MODELS=20 --OLLAMA_NUM_PARALLEL=20
-ENV OLLAMA_HOST "0.0.0.0:7860"
+# ENV OLLAMA_MAX_LOADED_MODELS=20 --OLLAMA_NUM_PARALLEL=20
+# ENV OLLAMA_HOST "0.0.0.0:7860"
 
-CMD ["serve"]
-EXPOSE 7860
+# CMD ["serve"]
+# EXPOSE 7860
+
+FROM python:3
+
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+COPY --chown=user . $HOME/LLM_API
+
+WORKDIR $HOME/LLM_API
+
+RUN mkdir $HOME/.cache
+
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+CMD ["gunicorn", "-w", "5", "-b", "0.0.0.0:7860","main:app"]

main.py

@@ -1,24 +1,39 @@
 from flask import Flask
 from flask import request
-from langchain_community.llms import Ollama
+from groq import Groq
+
+import os
 
 app = Flask(__name__)
+client = Groq(
+    api_key=os.environ.get("GROQ_API_KEY")
+)
 
-llm = Ollama(model="phi3")
 
-@app.route('/completion', methods=['POST'])
+@app.route("/api/generate", methods=['POST'])
 def completion():
     """
     {
-        user: <username>,
-        text: <text>
+        "model": "llama3-70b-8192",
+        "prompt": "why is the sky blue?"
     }
     """
-    
+
     message = request.get_json()
-    llm_output = llm.invoke(message['text'])
     
-    print(llm_output)
-    return llm_output
+    model = message['model']
+    prompt = message['prompt']
+
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": prompt,
+            }
+        ],
+        model=model,
+    )
+
+    return chat_completion.choices[0].message.content 
 
-# curl -v -X POST 'http://127.0.0.1:8000/completion' --header 'Content-Type: application/json' --data '{"user": "test-user", "text": "What do you know about 3D graphics"}'
+# curl -v -X POST 'http://127.0.0.1:8000/api/generate' --header 'Content-Type: application/json' --data '{"model": "llama3-70b-8192", "prompt": "why is sky blue?"}'

requirements.txt

@@ -9,6 +9,7 @@ charset-normalizer==3.3.2
 click==8.1.7
 dataclasses-json==0.6.6
 diskcache==5.6.3
+distro==1.9.0
 dnspython==2.6.1
 email_validator==2.1.1
 fastapi==0.111.0
@@ -18,6 +19,7 @@ Flask==3.0.3
 frozenlist==1.4.1
 fsspec==2024.5.0
 greenlet==3.0.3
+groq==0.8.0
 gunicorn==22.0.0
 h11==0.14.0
 httpcore==1.0.5
@@ -58,6 +60,7 @@ referencing==0.35.1
 requests==2.31.0
 rich==13.7.1
 rpds-py==0.18.1
+setuptools==70.0.0
 shellingham==1.5.4
 sniffio==1.3.1
 SQLAlchemy==2.0.30
@@ -76,4 +79,5 @@ uvloop==0.19.0
 watchfiles==0.21.0
 websockets==12.0
 Werkzeug==3.0.3
+wheel==0.43.0
 yarl==1.9.4