Upload model_config.py with huggingface_hub

model_config.py

@@ -0,0 +1,208 @@
+import os
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import logging
+
+logger = logging.getLogger(__name__)
+
+class ModelConfig:
+    """Configuration for different LLM models optimized for Hugging Face Spaces"""
+    
+    MODELS = {
+        "dialogpt-medium": {
+            "name": "microsoft/DialoGPT-medium",
+            "description": "Conversational AI model, good for chat",
+            "max_length": 512,
+            "memory_usage": "medium",
+            "recommended_for": "chat, conversation"
+        },
+        "dialogpt-small": {
+            "name": "microsoft/DialoGPT-small", 
+            "description": "Smaller conversational model, faster inference",
+            "max_length": 256,
+            "memory_usage": "low",
+            "recommended_for": "quick responses, limited resources"
+        },
+        "gpt2": {
+            "name": "gpt2",
+            "description": "General purpose text generation",
+            "max_length": 1024,
+            "memory_usage": "medium",
+            "recommended_for": "text generation, creative writing"
+        },
+        "distilgpt2": {
+            "name": "distilgpt2",
+            "description": "Distilled GPT-2, faster and smaller",
+            "max_length": 512,
+            "memory_usage": "low", 
+            "recommended_for": "fast inference, resource constrained"
+        },
+        "flan-t5-small": {
+            "name": "google/flan-t5-small",
+            "description": "Instruction-tuned T5 model",
+            "max_length": 512,
+            "memory_usage": "low",
+            "recommended_for": "instruction following, Q&A"
+        }
+    }
+    
+    @classmethod
+    def get_model_info(cls, model_key: str = None):
+        """Get information about available models"""
+        if model_key:
+            return cls.MODELS.get(model_key)
+        return cls.MODELS
+    
+    @classmethod
+    def get_recommended_model(cls, use_case: str = "general"):
+        """Get recommended model based on use case"""
+        recommendations = {
+            "chat": "dialogpt-medium",
+            "fast": "distilgpt2", 
+            "general": "gpt2",
+            "qa": "flan-t5-small",
+            "low_memory": "dialogpt-small"
+        }
+        return recommendations.get(use_case, "dialogpt-medium")
+
+class ModelManager:
+    """Manages model loading and inference"""
+    
+    def __init__(self, model_name: str = None):
+        self.model_name = model_name or os.getenv("MODEL_NAME", "microsoft/DialoGPT-medium")
+        self.model = None
+        self.tokenizer = None
+        self.pipeline = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.loaded = False
+        
+    def load_model(self):
+        """Load the specified model"""
+        try:
+            logger.info(f"Loading model: {self.model_name}")
+            logger.info(f"Using device: {self.device}")
+            
+            # Load tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_name,
+                padding_side="left"
+            )
+            
+            # Add padding token if it doesn't exist
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            # Load model with optimizations
+            model_kwargs = {
+                "low_cpu_mem_usage": True,
+                "torch_dtype": torch.float16 if self.device == "cuda" else torch.float32,
+            }
+            
+            if self.device == "cuda":
+                model_kwargs["device_map"] = "auto"
+            
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                **model_kwargs
+            )
+            
+            # Move to device if not using device_map
+            if self.device == "cpu":
+                self.model = self.model.to(self.device)
+            
+            # Create pipeline
+            self.pipeline = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                device=0 if self.device == "cuda" else -1,
+                return_full_text=False
+            )
+            
+            self.loaded = True
+            logger.info("Model loaded successfully!")
+            
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise e
+    
+    def generate_response(self, 
+                         prompt: str, 
+                         max_length: int = 100,
+                         temperature: float = 0.7,
+                         top_p: float = 0.9,
+                         do_sample: bool = True) -> str:
+        """Generate response using the loaded model"""
+        
+        if not self.loaded:
+            raise RuntimeError("Model not loaded. Call load_model() first.")
+        
+        try:
+            # Generate response
+            outputs = self.pipeline(
+                prompt,
+                max_new_tokens=max_length,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=do_sample,
+                pad_token_id=self.tokenizer.eos_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                truncation=True
+            )
+            
+            # Extract generated text
+            if outputs and len(outputs) > 0:
+                generated_text = outputs[0]['generated_text']
+                return generated_text.strip()
+            else:
+                return "Sorry, I couldn't generate a response."
+                
+        except Exception as e:
+            logger.error(f"Error generating response: {str(e)}")
+            raise e
+    
+    def get_model_info(self):
+        """Get information about the loaded model"""
+        return {
+            "model_name": self.model_name,
+            "device": self.device,
+            "loaded": self.loaded,
+            "tokenizer_vocab_size": len(self.tokenizer) if self.tokenizer else None,
+            "model_parameters": sum(p.numel() for p in self.model.parameters()) if self.model else None
+        }
+    
+    def unload_model(self):
+        """Unload the model to free memory"""
+        if self.model:
+            del self.model
+            self.model = None
+        if self.tokenizer:
+            del self.tokenizer
+            self.tokenizer = None
+        if self.pipeline:
+            del self.pipeline
+            self.pipeline = None
+        
+        # Clear CUDA cache if using GPU
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
+        self.loaded = False
+        logger.info("Model unloaded successfully")
+
+# Global model manager instance
+model_manager = None
+
+def get_model_manager(model_name: str = None) -> ModelManager:
+    """Get or create the global model manager instance"""
+    global model_manager
+    if model_manager is None:
+        model_manager = ModelManager(model_name)
+    return model_manager
+
+def initialize_model(model_name: str = None):
+    """Initialize and load the model"""
+    manager = get_model_manager(model_name)
+    if not manager.loaded:
+        manager.load_model()
+    return manager