Spaces:
Running
Running
| """ | |
| Hugging Face service for model inference and API integration. | |
| """ | |
| import os | |
| import logging | |
| from typing import Dict, List, Optional, Any | |
| # Optional imports - handle missing dependencies gracefully | |
| try: | |
| from huggingface_hub import InferenceClient | |
| HUGGINGFACE_HUB_AVAILABLE = True | |
| except ImportError: | |
| HUGGINGFACE_HUB_AVAILABLE = False | |
| InferenceClient = None | |
| try: | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
| TRANSFORMERS_AVAILABLE = True | |
| except ImportError: | |
| TRANSFORMERS_AVAILABLE = False | |
| AutoTokenizer = None | |
| AutoModelForCausalLM = None | |
| pipeline = None | |
| try: | |
| import torch | |
| TORCH_AVAILABLE = True | |
| except ImportError: | |
| TORCH_AVAILABLE = False | |
| torch = None | |
| from config import settings | |
| logger = logging.getLogger(__name__) | |
| class HuggingFaceService: | |
| """Service for Hugging Face model interactions.""" | |
| def __init__(self): | |
| self.client = None | |
| self.tokenizer = None | |
| self.model = None | |
| self.pipeline = None | |
| # Check if dependencies are available | |
| if not HUGGINGFACE_HUB_AVAILABLE: | |
| logger.warning("huggingface_hub not available. Hugging Face API features will be disabled.") | |
| if not TRANSFORMERS_AVAILABLE: | |
| logger.warning("transformers not available. Local model features will be disabled.") | |
| if not TORCH_AVAILABLE: | |
| logger.warning("torch not available. Local model features will be disabled.") | |
| # Initialize Hugging Face client if API token is provided and dependencies are available | |
| if settings.hf_api_token and HUGGINGFACE_HUB_AVAILABLE: | |
| try: | |
| self.client = InferenceClient(token=settings.hf_api_token) | |
| logger.info("Hugging Face Inference Client initialized successfully") | |
| except Exception as e: | |
| logger.error(f"Failed to initialize Hugging Face client: {e}") | |
| # Load local model if specified and dependencies are available | |
| if settings.hf_model_name and TRANSFORMERS_AVAILABLE and TORCH_AVAILABLE: | |
| self._load_local_model() | |
| def _load_local_model(self): | |
| """Load a local Hugging Face model.""" | |
| if not TRANSFORMERS_AVAILABLE or not TORCH_AVAILABLE: | |
| logger.error("Cannot load local model: transformers or torch not available") | |
| return | |
| try: | |
| logger.info(f"Loading local model: {settings.hf_model_name}") | |
| # Load tokenizer and model | |
| self.tokenizer = AutoTokenizer.from_pretrained(settings.hf_model_name) | |
| self.model = AutoModelForCausalLM.from_pretrained( | |
| settings.hf_model_name, | |
| torch_dtype=torch.float16, | |
| device_map="auto" | |
| ) | |
| # Create pipeline | |
| self.pipeline = pipeline( | |
| "text-generation", | |
| model=self.model, | |
| tokenizer=self.tokenizer, | |
| max_length=2048, | |
| temperature=0.7, | |
| do_sample=True | |
| ) | |
| logger.info("Local model loaded successfully") | |
| except Exception as e: | |
| logger.error(f"Failed to load local model: {e}") | |
| async def generate_text( | |
| self, | |
| prompt: str, | |
| model_name: Optional[str] = None, | |
| max_length: int = 2048, | |
| temperature: float = 0.7, | |
| use_local: bool = False | |
| ) -> str: | |
| """Generate text using Hugging Face models.""" | |
| try: | |
| if use_local and self.pipeline and TRANSFORMERS_AVAILABLE: | |
| # Use local model | |
| result = self.pipeline( | |
| prompt, | |
| max_length=max_length, | |
| temperature=temperature, | |
| do_sample=True, | |
| pad_token_id=self.tokenizer.eos_token_id | |
| ) | |
| return result[0]['generated_text'] | |
| elif self.client and model_name and HUGGINGFACE_HUB_AVAILABLE: | |
| # Use Hugging Face Inference API | |
| result = self.client.text_generation( | |
| prompt, | |
| model=model_name, | |
| max_new_tokens=max_length, | |
| temperature=temperature, | |
| do_sample=True | |
| ) | |
| return result | |
| else: | |
| raise Exception("No model available. Please configure HF_API_TOKEN or HF_MODEL_NAME, or install required dependencies") | |
| except Exception as e: | |
| logger.error(f"Hugging Face text generation error: {e}") | |
| raise | |
| async def create_embedding( | |
| self, | |
| text: str, | |
| model_name: str = "sentence-transformers/all-MiniLM-L6-v2" | |
| ) -> List[float]: | |
| """Create embedding using Hugging Face models.""" | |
| try: | |
| if self.client and HUGGINGFACE_HUB_AVAILABLE: | |
| # Use Hugging Face Inference API | |
| result = self.client.feature_extraction( | |
| text, | |
| model=model_name | |
| ) | |
| return result[0] if isinstance(result, list) else result | |
| else: | |
| raise Exception("Hugging Face client not initialized or dependencies not available") | |
| except Exception as e: | |
| logger.error(f"Hugging Face embedding error: {e}") | |
| raise | |
| async def classify_text( | |
| self, | |
| text: str, | |
| model_name: str = "distilbert-base-uncased-finetuned-sst-2-english" | |
| ) -> Dict[str, Any]: | |
| """Classify text using Hugging Face models.""" | |
| try: | |
| if self.client and HUGGINGFACE_HUB_AVAILABLE: | |
| result = self.client.text_classification( | |
| text, | |
| model=model_name | |
| ) | |
| return result | |
| else: | |
| raise Exception("Hugging Face client not initialized or dependencies not available") | |
| except Exception as e: | |
| logger.error(f"Hugging Face classification error: {e}") | |
| raise | |
| async def translate_text( | |
| self, | |
| text: str, | |
| source_lang: str = "en", | |
| target_lang: str = "es", | |
| model_name: str = "Helsinki-NLP/opus-mt-en-es" | |
| ) -> str: | |
| """Translate text using Hugging Face models.""" | |
| try: | |
| if self.client and HUGGINGFACE_HUB_AVAILABLE: | |
| result = self.client.translation( | |
| text, | |
| model=model_name | |
| ) | |
| return result | |
| else: | |
| raise Exception("Hugging Face client not initialized or dependencies not available") | |
| except Exception as e: | |
| logger.error(f"Hugging Face translation error: {e}") | |
| raise | |
| def get_available_models(self) -> List[str]: | |
| """Get list of available models.""" | |
| models = [] | |
| if self.client and HUGGINGFACE_HUB_AVAILABLE: | |
| models.append("Hugging Face Inference API (remote)") | |
| if self.pipeline and TRANSFORMERS_AVAILABLE: | |
| models.append(f"Local model: {settings.hf_model_name}") | |
| if not models: | |
| models.append("No models available - check dependencies and configuration") | |
| return models | |
| def get_service_status(self) -> Dict[str, Any]: | |
| """Get service status and configuration.""" | |
| return { | |
| "service": "huggingface", | |
| "status": "available" if (self.client or self.pipeline) else "unavailable", | |
| "dependencies": { | |
| "huggingface_hub": HUGGINGFACE_HUB_AVAILABLE, | |
| "transformers": TRANSFORMERS_AVAILABLE, | |
| "torch": TORCH_AVAILABLE | |
| }, | |
| "client_initialized": self.client is not None, | |
| "local_model_loaded": self.pipeline is not None, | |
| "api_token_configured": bool(settings.hf_api_token), | |
| "local_model_configured": bool(settings.hf_model_name), | |
| "available_models": self.get_available_models() | |
| } |