fullstuckdev/medication-ai-model
1
1import os2from fastapi import FastAPI, HTTPException, BackgroundTasks3from fastapi.middleware.cors import CORSMiddleware4from pydantic import BaseModel5import torch6from transformers import AutoTokenizer, AutoModelForCausalLM7import logging8from typing import List, Optional9from datasets import load_dataset10from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling11import json12 13# Setup logging14logging.basicConfig(level=logging.INFO)15logger = logging.getLogger(__name__)16 17# Setup cache directory18os.makedirs("/app/cache", exist_ok=True)19os.environ['TRANSFORMERS_CACHE'] = "/app/cache"20 21# Pydantic models for request/response22class GenerateRequest(BaseModel):23 text: str24 max_length: Optional[int] = 51225 temperature: Optional[float] = 0.726 num_return_sequences: Optional[int] = 127 28class GenerateResponse(BaseModel):29 generated_text: List[str]30 31class HealthResponse(BaseModel):32 status: str33 model_loaded: bool34 gpu_available: bool35 device: str36 37class TrainRequest(BaseModel):38 dataset_path: str39 num_epochs: Optional[int] = 340 batch_size: Optional[int] = 441 learning_rate: Optional[float] = 2e-542 43class TrainResponse(BaseModel):44 status: str45 message: str46 47# Add training status tracking48class TrainingStatus:49 def __init__(self):50 self.is_training = False51 self.current_epoch = 052 self.current_loss = None53 self.status = "idle"54 55training_status = TrainingStatus()56 57# Initialize FastAPI app58app = FastAPI(59 title="Medical LLaMA API",60 description="API for medical text generation using fine-tuned LLaMA model",61 version="1.0.0",62 docs_url="/docs",63 redoc_url="/redoc"64)65 66# Add CORS middleware67app.add_middleware(68 CORSMiddleware,69 allow_origins=["*"],70 allow_credentials=True,71 allow_methods=["*"],72 allow_headers=["*"],73)74 75# Global variables for model and tokenizer76model = None77tokenizer = None78 79@app.get("/", response_model=HealthResponse, tags=["Health"])80async def root():81 """82 Root endpoint to check API health and model status83 """84 device = "cuda" if torch.cuda.is_available() else "cpu"85 return HealthResponse(86 status="online",87 model_loaded=model is not None,88 gpu_available=torch.cuda.is_available(),89 device=device90 )91 92@app.post("/generate", response_model=GenerateResponse, tags=["Generation"])93async def generate_text(request: GenerateRequest):94 """95 Generate medical text based on input prompt96 """97 try:98 # Check if model is loaded99 if model is None or tokenizer is None:100 logger.error("Model or tokenizer not initialized")101 raise HTTPException(102 status_code=500, 103 detail="Model not loaded. Please check if model was initialized correctly."104 )105 106 logger.info(f"Generating text for input: {request.text[:50]}...")107 108 # Log device information109 device_info = f"Using device: {model.device}"110 logger.info(device_info)111 112 # Tokenize input113 try:114 inputs = tokenizer(115 request.text,116 return_tensors="pt",117 padding=True,118 truncation=True,119 max_length=request.max_length120 )121 logger.info("Input tokenized successfully")122 123 # Move inputs to correct device124 inputs = {k: v.to(model.device) for k, v in inputs.items()}125 126 except Exception as e:127 logger.error(f"Tokenization error: {str(e)}")128 raise HTTPException(status_code=500, detail=f"Tokenization failed: {str(e)}")129 130 # Generate text131 try:132 with torch.no_grad():133 generated_ids = model.generate(134 inputs.input_ids,135 max_length=request.max_length,136 num_return_sequences=request.num_return_sequences,137 temperature=request.temperature,138 pad_token_id=tokenizer.pad_token_id,139 eos_token_id=tokenizer.eos_token_id,140 )141 logger.info("Text generated successfully")142 except Exception as e:143 logger.error(f"Generation error: {str(e)}")144 raise HTTPException(status_code=500, detail=f"Text generation failed: {str(e)}")145 146 # Decode generated text147 try:148 generated_texts = [149 tokenizer.decode(g, skip_special_tokens=True)150 for g in generated_ids151 ]152 logger.info("Text decoded successfully")153 except Exception as e:154 logger.error(f"Decoding error: {str(e)}")155 raise HTTPException(status_code=500, detail=f"Text decoding failed: {str(e)}")156 157 return GenerateResponse(generated_text=generated_texts)158 159 except HTTPException as he:160 raise he161 except Exception as e:162 logger.error(f"Unexpected error: {str(e)}")163 raise HTTPException(164 status_code=500, 165 detail=f"An unexpected error occurred: {str(e)}"166 )167 168@app.get("/health", tags=["Health"])169async def health_check():170 """171 Check the health status of the API and model172 """173 return {174 "status": "healthy",175 "model_loaded": model is not None,176 "gpu_available": torch.cuda.is_available(),177 "device": "cuda" if torch.cuda.is_available() else "cpu"178 }179 180@app.on_event("startup")181async def startup_event():182 logger.info("Starting up application...")183 try:184 global tokenizer, model185 tokenizer, model = init_model()186 logger.info("Model loaded successfully")187 except Exception as e:188 logger.error(f"Failed to load model: {str(e)}")189 190@app.post("/train", response_model=TrainResponse, tags=["Training"])191async def train_model(request: TrainRequest, background_tasks: BackgroundTasks):192 """193 Start model training with the specified dataset194 195 Parameters:196 - dataset_path: Path to the JSON dataset file197 - num_epochs: Number of training epochs198 - batch_size: Training batch size199 - learning_rate: Learning rate for training200 """201 if training_status.is_training:202 raise HTTPException(status_code=400, detail="Training is already in progress")203 204 try:205 # Verify dataset exists206 if not os.path.exists(request.dataset_path):207 raise HTTPException(status_code=404, detail="Dataset file not found")208 209 # Start training in background210 background_tasks.add_task(211 run_training,212 request.dataset_path,213 request.num_epochs,214 request.batch_size,215 request.learning_rate216 )217 218 return TrainResponse(219 status="started",220 message="Training started in background"221 )222 223 except Exception as e:224 logger.error(f"Training setup error: {str(e)}")225 raise HTTPException(status_code=500, detail=str(e))226 227@app.get("/train/status", tags=["Training"])228async def get_training_status():229 """230 Get current training status231 """232 return {233 "is_training": training_status.is_training,234 "current_epoch": training_status.current_epoch,235 "current_loss": training_status.current_loss,236 "status": training_status.status237 }238 239# Add training function240async def run_training(dataset_path: str, num_epochs: int, batch_size: int, learning_rate: float):241 global model, tokenizer, training_status242 243 try:244 training_status.is_training = True245 training_status.status = "loading_dataset"246 247 # Load dataset248 dataset = load_dataset("json", data_files=dataset_path)249 250 training_status.status = "preprocessing"251 252 # Preprocess function253 def preprocess_function(examples):254 return tokenizer(255 examples["text"],256 truncation=True,257 padding="max_length",258 max_length=512259 )260 261 # Tokenize dataset262 tokenized_dataset = dataset.map(263 preprocess_function,264 batched=True,265 remove_columns=dataset["train"].column_names266 )267 268 training_status.status = "training"269 270 # Training arguments271 training_args = TrainingArguments(272 output_dir=f"{model_output_path}/checkpoints",273 per_device_train_batch_size=batch_size,274 gradient_accumulation_steps=4,275 num_train_epochs=num_epochs,276 learning_rate=learning_rate,277 fp16=True,278 save_steps=500,279 logging_steps=100,280 )281 282 # Initialize trainer283 trainer = Trainer(284 model=model,285 args=training_args,286 train_dataset=tokenized_dataset["train"],287 data_collator=DataCollatorForLanguageModeling(288 tokenizer=tokenizer,289 mlm=False290 ),291 )292 293 # Training callback to update status294 class TrainingCallback(trainer.callback_handler):295 def on_epoch_begin(self, args, state, control, **kwargs):296 training_status.current_epoch = state.epoch297 298 def on_log(self, args, state, control, logs=None, **kwargs):299 if logs:300 training_status.current_loss = logs.get("loss", None)301 302 trainer.add_callback(TrainingCallback)303 304 # Start training305 trainer.train()306 307 # Save the model308 training_status.status = "saving"309 model.save_pretrained(model_output_path)310 tokenizer.save_pretrained(model_output_path)311 312 training_status.status = "completed"313 logger.info("Training completed successfully")314 315 except Exception as e:316 training_status.status = f"failed: {str(e)}"317 logger.error(f"Training error: {str(e)}")318 raise319 320 finally:321 training_status.is_training = False322 323# Update model initialization324def init_model():325 try:326 device = "cuda" if torch.cuda.is_available() else "cpu"327 logger.info(f"Loading model on device: {device}")328 329 model_name = "nvidia/Meta-Llama-3.2-3B-Instruct-ONNX-INT4"330 331 # Load tokenizer332 logger.info("Loading tokenizer...")333 tokenizer = AutoTokenizer.from_pretrained(334 model_name,335 cache_dir="/app/cache",336 trust_remote_code=True337 )338 339 # Add padding token if not present340 if tokenizer.pad_token is None:341 tokenizer.pad_token = tokenizer.eos_token342 343 logger.info("Loading model...")344 model = AutoModelForCausalLM.from_pretrained(345 model_name,346 torch_dtype=torch.float16 if device == "cuda" else torch.float32,347 device_map="auto",348 cache_dir="/app/cache",349 trust_remote_code=True350 )351 352 logger.info(f"Model loaded successfully on {device}")353 return tokenizer, model354 355 except Exception as e:356 logger.error(f"Model initialization error: {str(e)}")357 raise358 359@app.get("/model-status", tags=["Health"])360async def model_status():361 """362 Get detailed model status363 """364 try:365 model_info = {366 "model_loaded": model is not None,367 "tokenizer_loaded": tokenizer is not None,368 "model_device": str(model.device) if model else None,369 "gpu_available": torch.cuda.is_available(),370 "cuda_device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0,371 "cuda_device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,372 "model_type": type(model).__name__ if model else None,373 "tokenizer_type": type(tokenizer).__name__ if tokenizer else None,374 }375 376 if model is not None:377 try:378 # Test tokenizer379 test_input = tokenizer("test", return_tensors="pt")380 model_info["tokenizer_working"] = True381 except Exception as e:382 model_info["tokenizer_working"] = False383 model_info["tokenizer_error"] = str(e)384 385 try:386 # Test model forward pass387 with torch.no_grad():388 test_output = model.generate(389 test_input.input_ids.to(model.device),390 max_length=10391 )392 model_info["model_working"] = True393 except Exception as e:394 model_info["model_working"] = False395 model_info["model_error"] = str(e)396 397 return model_info398 399 except Exception as e:400 logger.error(f"Error checking model status: {str(e)}")401 return {402 "error": str(e),403 "model_loaded": model is not None,404 "tokenizer_loaded": tokenizer is not None405 }