feat: добавлена загрузка модели в папку КЭШа

main
Artem-Darius Weber 2 months ago
parent 219ad09569
commit db02ddbc62

@ -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")

@ -4,6 +4,7 @@
# #
import os
import asyncio import asyncio
import json import json
from aiokafka import AIOKafkaConsumer, AIOKafkaProducer from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
@ -11,11 +12,13 @@ from datetime import datetime
import torch import torch
from transformers import AutoTokenizer, AutoModelForCausalLM from transformers import AutoTokenizer, AutoModelForCausalLM
cache_dir = "./models"
model_id = "NousResearch/Meta-Llama-3.1-8B" model_id = "NousResearch/Meta-Llama-3.1-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id) tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
model_id, model_id,
cache_dir=cache_dir,
torch_dtype=torch.float32, # Использование float32 для CPU torch_dtype=torch.float32, # Использование float32 для CPU
device_map=None # Явное указание, что модель не будет использовать GPU device_map=None # Явное указание, что модель не будет использовать GPU
) )

Loading…
Cancel
Save