From 6cd4a15b5f745607373bd6316c6348ca164ed50a Mon Sep 17 00:00:00 2001
From: Viet-Anh Nguyen <vietanh.dev@gmail.com>
Date: Fri, 27 Sep 2024 23:17:38 +0700
Subject: [PATCH] Update workflow

---
 .github/workflows/publish-to-pypi.yml | 62 +++++++++++++++++++--------
 llama_assistant/llama_assistant.py    | 27 ++++++------
 pyproject.toml                        |  2 +-
 requirements.txt                      |  4 +-
 4 files changed, 62 insertions(+), 33 deletions(-)

diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml
index a43672f..ca48535 100644
--- a/.github/workflows/publish-to-pypi.yml
+++ b/.github/workflows/publish-to-pypi.yml
@@ -5,23 +5,51 @@ on:
     tags:
       - 'v*'  # Trigger on tags starting with 'v'
 
+permissions:
+  contents: read
+
 jobs:
-  publish:
+  release-build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0  # Fetch all history for all tags and branches
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"  # Specify the minimum required Python version
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+          pip install build setuptools_scm
+
+      - name: Build release distributions
+        run: python -m build
+
+      - name: Upload distributions
+        uses: actions/upload-artifact@v4
+        with:
+          name: release-dists
+          path: dist/
+
+  pypi-publish:
     runs-on: ubuntu-latest
+    needs:
+      - release-build
+    permissions:
+      id-token: write
+    environment:
+      name: release
     steps:
-    - uses: actions/checkout@v3
-    - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.x'
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install setuptools wheel twine
-    - name: Build package
-      run: python setup.py sdist bdist_wheel
-    - name: Publish to PyPI
-      env:
-        TWINE_USERNAME: __token__
-        TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
-      run: twine upload dist/*
\ No newline at end of file
+      - name: Retrieve release distributions
+        uses: actions/download-artifact@v4
+        with:
+          name: release-dists
+          path: dist/
+
+      - name: Publish release distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
\ No newline at end of file
diff --git a/llama_assistant/llama_assistant.py b/llama_assistant/llama_assistant.py
index 07bbf3d..7fd5e09 100644
--- a/llama_assistant/llama_assistant.py
+++ b/llama_assistant/llama_assistant.py
@@ -1,6 +1,6 @@
 import json
 import markdown
-from mlx_lm import load, generate
+from llama_cpp import Llama
 from pathlib import Path
 
 from PyQt6.QtWidgets import (
@@ -59,7 +59,10 @@ def __init__(self):
         self.image_label = None
 
     def load_model(self):
-        self.model, self.tokenizer = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
+        self.model = Llama.from_pretrained(
+            repo_id="hugging-quants/Llama-3.2-1B-Instruct-Q4_K_M-GGUF",
+            filename="*q4_k_m.gguf",
+        )
 
     def load_settings(self):
         home_dir = Path.home()
@@ -350,18 +353,16 @@ def process_text(self, message, task="chat"):
         elif task == "write email":
             prompt = f"Write an email about: {message}"
 
-        if (
-            hasattr(self.tokenizer, "apply_chat_template")
-            and self.tokenizer.chat_template is not None
-        ):
-            messages = [{"role": "user", "content": prompt}]
-            prompt = self.tokenizer.apply_chat_template(
-                messages, tokenize=False, add_generation_prompt=True
-            )
-
-        response = generate(
-            self.model, self.tokenizer, prompt=prompt, max_tokens=2048, verbose=False
+        output = self.model.create_chat_completion(
+            messages = [
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ]
         )
+        response = output["choices"][0]["message"]["content"]
+
         self.last_response = response
 
         self.chat_box.append(f"<b>You:</b> {message}")
diff --git a/pyproject.toml b/pyproject.toml
index 95e439a..6f49952 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,7 @@ classifiers = [
 dependencies = [
     "PyQt6",
     "markdown",
-    "mlx_lm",
+    "llama-cpp-python",
     "pynput",
     "SpeechRecognition",
 ]
diff --git a/requirements.txt b/requirements.txt
index f7c8023..542f13d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 PyQt6==6.7.1
 SpeechRecognition==3.10.4
 markdown==3.7
-mlx-lm==0.18.2
-pynput==1.7.7
\ No newline at end of file
+pynput==1.7.7
+llama-cpp-python
\ No newline at end of file