Spaces:
Running
Running
File size: 7,202 Bytes
8115389 |
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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
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
|