From db02ddbc62c3dbf39be4a2e2dc987e16cf4d1657 Mon Sep 17 00:00:00 2001
From: Artem Darius Weber <mit.3tlasa@gmail.com>
Date: Sun, 10 Nov 2024 01:39:00 +0300
Subject: [PATCH] =?UTF-8?q?feat:=20=D0=B4=D0=BE=D0=B1=D0=B0=D0=B2=D0=BB?=
 =?UTF-8?q?=D0=B5=D0=BD=D0=B0=20=D0=B7=D0=B0=D0=B3=D1=80=D1=83=D0=B7=D0=BA?=
 =?UTF-8?q?=D0=B0=20=D0=BC=D0=BE=D0=B4=D0=B5=D0=BB=D0=B8=20=D0=B2=20=D0=BF?=
 =?UTF-8?q?=D0=B0=D0=BF=D0=BA=D1=83=20=D0=9A=D0=AD=D0=A8=D0=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 agentic/llm/_download_model.py | 16 ++++++++++++++++
 agentic/llm/llm_worker.py      |  3 +++
 2 files changed, 19 insertions(+)
 create mode 100644 agentic/llm/_download_model.py

diff --git a/agentic/llm/_download_model.py b/agentic/llm/_download_model.py
new file mode 100644
index 0000000..457771d
--- /dev/null
+++ b/agentic/llm/_download_model.py
@@ -0,0 +1,16 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
+
+cache_dir = "./models"
+model_id = "NousResearch/Meta-Llama-3.1-8B"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    cache_dir=cache_dir,
+    torch_dtype=torch.float32,   # Использование float32 для CPU
+    device_map=None              # Явное указание, что модель не будет использовать GPU
+)
+
+model.to("cpu")
diff --git a/agentic/llm/llm_worker.py b/agentic/llm/llm_worker.py
index 619ea9a..dad5bcc 100644
--- a/agentic/llm/llm_worker.py
+++ b/agentic/llm/llm_worker.py
@@ -4,6 +4,7 @@
 #
 
 
+import os
 import asyncio
 import json
 from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
@@ -11,11 +12,13 @@ from datetime import datetime
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+cache_dir = "./models"
 model_id = "NousResearch/Meta-Llama-3.1-8B"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
+    cache_dir=cache_dir,
     torch_dtype=torch.float32,   # Использование float32 для CPU
     device_map=None              # Явное указание, что модель не будет использовать GPU
 )