Spaces:
Sleeping
Sleeping
File size: 3,173 Bytes
06cebb7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoModelForSequenceClassification
import sys
from pathlib import Path
try:
from utils.reproducibility import set_seed
except ModuleNotFoundError:
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from utils.reproducibility import set_seed
# Establecer una semilla para reproducibilidad (descomenta para activar)
# set_seed(72)
# Clase encargada de cargar y gestionar los modelos y tokenizadores utilizados en el proyecto.
class ModelManager:
def __init__(self, config):
self.config = config
self.models = {}
self.tokenizers = {}
# Carga el modelo generador de prompts y su tokenizer con cuantización 4-bit (bitsandbytes)
def load_generator(self):
modelo_id = self.config.modelo_generador['id_modelo']
modo = self.config.modelo_generador['modo_interaccion']
if modelo_id != '' and str(modo).lower() == 'local':
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype="float16",
bnb_4bit_quant_type="nf4"
)
tokenizer = AutoTokenizer.from_pretrained(modelo_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(modelo_id, quantization_config=bnb_config, low_cpu_mem_usage=True, device_map="auto")
self.models['generator'] = model
self.tokenizers['generator'] = tokenizer
# Carga el modelo LLM a evaluar y su tokenizer con cuantización 4-bit.
def load_evaluator(self):
modelo_id = self.config.modelo_a_evaluar['id_modelo']
modo = self.config.modelo_a_evaluar['modo_interaccion']
if modelo_id != '' and str(modo).lower() == 'local':
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype="float16",
bnb_4bit_quant_type="nf4"
)
tokenizer = AutoTokenizer.from_pretrained(modelo_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(modelo_id, quantization_config=bnb_config, low_cpu_mem_usage=True, device_map="auto")
self.models['evaluator'] = model
self.tokenizers['evaluator'] = tokenizer
# Carga el modelo de análisis de sentimiento (SequenceClassification) en CUDA y lo pone en eval().
def load_sentiment(self):
modelo_id = self.config.modelo_analisis_de_sentimiento['id_modelo']
modo = self.config.modelo_analisis_de_sentimiento['modo_interaccion']
if modelo_id != '' and str(modo).lower() == 'local':
tokenizer = AutoTokenizer.from_pretrained(modelo_id, use_fast=True)
model = AutoModelForSequenceClassification.from_pretrained(modelo_id, low_cpu_mem_usage=True)
model = model.to("cuda")
model.eval()
self.models['sentiment'] = model
self.tokenizers['sentiment'] = tokenizer
def get_model(self, key):
return self.models.get(key), self.tokenizers.get(key)
|