Upload inference.py with huggingface_hub

inference.py

@@ -1,73 +1,189 @@
 import google.generativeai as genai
-import os
-from dotenv import load_dotenv
-import yaml
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from huggingface_utils import HuggingFaceUtils
+import time
 
 class InferenceManager:
-    def __init__(self):
-        # Cargar configuración
-        with open('config.yaml', 'r') as file:
-            self.config = yaml.safe_load(file)
-        
-        # Configurar Gemini
-        load_dotenv()
-        genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
-        self.gemini = genai.GenerativeModel('gemini-pro')
-        
-        # Configurar Mixtral (se cargará bajo demanda)
-        self.mixtral = None
-        self.mixtral_tokenizer = None
-        
-        # Estado inicial
-        self.current_model = 'gemini'
-        self.current_mode = 'seguros'
-        
-    def load_mixtral(self):
-        if not self.mixtral:
-            self.mixtral_tokenizer = AutoTokenizer.from_pretrained(
-                self.config['models']['mixtral']['name']
-            )
-            self.mixtral = AutoModelForCausalLM.from_pretrained(
-                self.config['models']['mixtral']['name'],
-                torch_dtype=torch.float16,
-                device_map="auto"
+    def __init__(self, config):
+        """Inicializar modelos de IA con validación de configuración"""
+        try:
+            # Obtener API keys del diccionario anidado
+            api_keys = config.get('api_keys', {})
+            
+            # Configurar parámetros del modelo primero
+            self.generation_config = {
+                "temperature": 0.9,
+                "top_p": 1,
+                "top_k": 1,
+                "max_output_tokens": 2048,
+            }
+            
+            self.safety_settings = [
+                {
+                    "category": "HARM_CATEGORY_HARASSMENT",
+                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+                },
+                {
+                    "category": "HARM_CATEGORY_HATE_SPEECH",
+                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+                },
+                {
+                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+                },
+                {
+                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+                },
+            ]
+            
+            # Validar API keys
+            self.google_api_key = api_keys.get('GOOGLE_API_KEY', '')
+            self.huggingface_token = api_keys.get('HUGGINGFACE_TOKEN', '')
+            
+            # Configurar HuggingFace primero como fallback
+            if self.huggingface_token:
+                self.huggingface = HuggingFaceUtils(token=self.huggingface_token)
+            else:
+                print("Advertencia: No se configuró HuggingFace Token")
+                self.huggingface = None
+            
+            # Intentar configurar Gemini
+            self.max_retries = 3
+            if self.google_api_key:
+                try:
+                    genai.configure(api_key=self.google_api_key)
+                    self.model = self._init_gemini_model()
+                    self.current_model = 'gemini'
+                except Exception as e:
+                    print(f"Error inicializando Gemini: {e}")
+                    if self.huggingface:
+                        print("Cambiando a Mixtral como fallback")
+                        self.current_model = 'mixtral'
+                    else:
+                        raise ValueError("No hay modelos disponibles")
+            else:
+                if self.huggingface:
+                    print("Usando Mixtral como modelo principal")
+                    self.current_model = 'mixtral'
+                else:
+                    raise ValueError("No hay modelos disponibles")
+            
+        except Exception as e:
+            print(f"Error en inicialización: {e}")
+            raise
+    
+    def _init_gemini_model(self, attempt=1):
+        """Inicializar modelo Gemini con reintentos"""
+        try:
+            return genai.GenerativeModel(
+                'gemini-pro',
+                generation_config=self.generation_config,
+                safety_settings=self.safety_settings
             )
+        except Exception as e:
+            if attempt < self.max_retries:
+                time.sleep(2 ** attempt)
+                return self._init_gemini_model(attempt + 1)
+            raise
     
     def change_model(self, model_name):
-        if model_name in ['gemini', 'mixtral']:
-            self.current_model = model_name
-            if model_name == 'mixtral':
-                self.load_mixtral()
-            return True
-        return False
+        """Cambiar modelo de IA con validación"""
+        try:
+            if model_name == 'gemini' and hasattr(self, 'model'):
+                self.current_model = 'gemini'
+                return True
+            elif model_name == 'mixtral' and self.huggingface:
+                self.current_model = 'mixtral'
+                return True
+            else:
+                print(f"Modelo solicitado '{model_name}' no está disponible.")
+                return False
+        except Exception as e:
+            print(f"Error cambiando modelo IA: {e}")
+            return False
     
-    def change_mode(self, mode):
-        if mode in self.config['modes']:
-            self.current_mode = mode
-            return True
-        return False
+    def _generate_gemini_response(self, prompt, attempt=1):
+        """Generar respuesta de Gemini con reintentos y fallback a Mixtral"""
+        try:
+            response = self.model.generate_content(prompt)
+            return response.text
+        except Exception as e:
+            if "safety" in str(e).lower():
+                return "Lo siento, no puedo responder a eso debido a restricciones de seguridad."
+            elif attempt < self.max_retries:
+                time.sleep(2 ** attempt)
+                return self._generate_gemini_response(prompt, attempt + 1)
+            else:
+                print(f"Error con Gemini, intentando Mixtral: {e}")
+                if self.huggingface:
+                    self.current_model = 'mixtral'
+                    return self.huggingface.generate_response(prompt)
+                return "Lo siento, el servicio no está respondiendo. Por favor, intenta de nuevo en unos momentos."
     
-    def get_response(self, message):
+    def _generate_mixtral_response(self, prompt, attempt=1):
+        """Generar respuesta de Mixtral con reintentos"""
+        try:
+            return self.huggingface.generate_response(prompt)
+        except Exception as e:
+            if attempt < self.max_retries:
+                print(f"Error con Mixtral (intento {attempt}), reintentando: {e}")
+                time.sleep(2 ** attempt)
+                return self._generate_mixtral_response(prompt, attempt + 1)
+            raise
+
+    def get_response(self, prompt, context='', history=None):
+        """Generar respuesta usando el modelo actual con manejo de errores mejorado"""
         try:
-            # Obtener contexto del modo actual
-            context = self.config['modes'][self.current_mode]['context']
-            prompt = f"{context}\n\nUsuario: {message}\nAsistente:"
+            if not prompt:
+                return "No se proporcionó un mensaje"
+            
+            # Construir el prompt completo
+            full_prompt = context + "\n\n" if context else ""
+            
+            # Agregar historial si existe
+            if history:
+                full_prompt += "Historial de la conversación:\n"
+                for msg in history:
+                    role = "Usuario" if msg["role"] == "user" else "Asistente"
+                    full_prompt += f"{role}: {msg['content']}\n"
             
-            if self.current_model == 'gemini':
-                response = self.gemini.generate_content(prompt)
-                return response.text
-            else:  # mixtral
-                inputs = self.mixtral_tokenizer(prompt, return_tensors="pt").to("cuda")
-                outputs = self.mixtral.generate(
-                    **inputs,
-                    max_length=self.config['models']['mixtral']['max_length'],
-                    temperature=self.config['models']['mixtral']['temperature'],
-                    pad_token_id=self.mixtral_tokenizer.eos_token_id
-                )
-                return self.mixtral_tokenizer.decode(outputs[0], skip_special_tokens=True)
+            full_prompt += f"\nUsuario: {prompt}"
+            
+            # Generar respuesta según el modelo
+            print(f"Usando modelo: {self.current_model}")
+            try:
+                if self.current_model == 'gemini':
+                    response = self._generate_gemini_response(full_prompt)
+                    if not response or len(response.strip()) < 10:
+                        raise Exception("Respuesta inválida de Gemini")
+                    return response
+                elif self.current_model == 'mixtral' and self.huggingface:
+                    response = self._generate_mixtral_response(full_prompt)
+                    if not response or len(response.strip()) < 10:
+                        raise Exception("Respuesta inválida de Mixtral")
+                    return response
+                else:
+                    return "Lo siento, no hay modelos disponibles en este momento."
+            except Exception as model_error:
+                print(f"Error con el modelo {self.current_model}: {model_error}")
+                
+                # Si Gemini falla, intentar con Mixtral
+                if self.current_model == 'gemini' and self.huggingface:
+                    print("Cambiando a Mixtral como fallback")
+                    self.current_model = 'mixtral'
+                    response = self._generate_mixtral_response(full_prompt)
+                    if response and len(response.strip()) >= 10:
+                        return response
+                # Si Mixtral falla y tenemos Gemini disponible, intentar con Gemini
+                elif self.current_model == 'mixtral' and hasattr(self, 'model'):
+                    print("Cambiando a Gemini como fallback")
+                    self.current_model = 'gemini'
+                    response = self._generate_gemini_response(full_prompt)
+                    if response and len(response.strip()) >= 10:
+                        return response
+                
+                raise Exception("No se pudo obtener una respuesta válida de ningún modelo")
                 
         except Exception as e:
-            print(f"Error en inferencia: {e}")
-            return "Lo siento, hubo un error al procesar tu mensaje."
+            print(f"Error general en get_response: {e}")
+            return "Lo siento, hubo un error inesperado. Por favor, intenta de nuevo."