Spaces:
Sleeping
Sleeping
File size: 3,567 Bytes
3ff2a4f | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """
MultiBackendLLMEngine – Moteur LLM local pour Qwen 7B (Apache 2.0)
Utilise transformers pour charger le modèle en 4-bit si disponible.
"""
import logging
import os
from typing import Optional
logger = logging.getLogger(__name__)
class MultiBackendLLMEngine:
"""Moteur LLM multi-backend (priorité au local Qwen)"""
def __init__(self, model_name: str = "Qwen/Qwen2.5-7B-Instruct"):
self.model_name = model_name
self.tokenizer = None
self.model = None
self.device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
self._load_model()
def _load_model(self):
"""Charge le modèle et le tokenizer."""
try:
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch
logger.info(f"Chargement du modèle {self.model_name} sur {self.device}...")
# Tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True)
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
# Quantization 4-bit si GPU disponible
if self.device == "cuda":
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
self.model = AutoModelForCausalLM.from_pretrained(
self.model_name,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True
)
else:
self.model = AutoModelForCausalLM.from_pretrained(
self.model_name,
device_map="auto",
trust_remote_code=True
)
logger.info(f"✅ Modèle {self.model_name} chargé avec succès")
except Exception as e:
logger.error(f"Erreur lors du chargement du modèle: {e}")
self.model = None
self.tokenizer = None
raise
def generate(self, prompt: str, max_new_tokens: int = 512, temperature: float = 0.7) -> str:
"""Génère une réponse à partir d'un prompt."""
if not self.model or not self.tokenizer:
return "Modèle non chargé."
try:
inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096)
inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
outputs = self.model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
do_sample=True,
top_p=0.95,
pad_token_id=self.tokenizer.eos_token_id
)
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
# Enlever le prompt de la réponse
if response.startswith(prompt):
response = response[len(prompt):].strip()
return response
except Exception as e:
logger.error(f"Erreur lors de la génération: {e}")
return f"Erreur de génération: {e}"
def route(self, query: str) -> dict:
"""Point d'entrée compatible avec le MoERouter."""
response = self.generate(query)
return {"response": response, "source": "Qwen-7B-local"}
|