Update app.py

app.py

@@ -1,14 +1,11 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from typing import Optional, List
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from transformers import BitsAndBytesConfig
+from typing import Optional
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel
 import torch
 import os
-from datetime import datetime
-import re
 
 # =========================
 # CONFIG
@@ -17,6 +14,9 @@ import re
 BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
 LORA_MODEL = "Delta0723/techmind-pro-v9"
 
+# Crear carpeta para offload si no existe
+os.makedirs("offload", exist_ok=True)
+
 # =========================
 # FastAPI Setup
 # =========================
@@ -40,24 +40,25 @@ try:
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)
     tokenizer.pad_token = tokenizer.eos_token
 
-quant_config = BitsAndBytesConfig(load_in_4bit=True)
+    quant_config = BitsAndBytesConfig(load_in_4bit=True)
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        device_map="auto",
+        trust_remote_code=True,
+        offload_folder="offload",
+        quantization_config=quant_config
+    )
 
-base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    device_map="auto",
-    trust_remote_code=True,
-    offload_folder="offload",
-    quantization_config=quant_config
-    
     model = PeftModel.from_pretrained(base_model, LORA_MODEL)
     model.eval()
 
+    print("✅ Modelo listo para usar")
+
 except Exception as e:
     print("❌ Error al cargar el modelo:", e)
     raise e
 
-print("✅ Modelo listo")
-
 # =========================
 # Data Models
 # =========================