Upload 3 files

Dockerfile

@@ -1,26 +1,24 @@
-FROM python:3.10
+FROM python:3.10 as builder
 
 WORKDIR /code
 
 COPY ./requirements.txt /code/requirements.txt
 
-# Upgrade pip before installing requirements
-RUN pip install --no-cache-dir --upgrade pip
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r /code/requirements.txt
 
-# Install dependencies
-RUN pip install --no-cache-dir -r /code/requirements.txt
+FROM python:3.10-slim
 
-# Create and use non-root user
 RUN useradd -m user
 USER user
 
-# Set environment variables
 ENV HOME=/home/user \
-    PATH=/home/user/.local/bin:${PATH}
+    PATH=/home/user/.local/bin:${PATH} \
+    PYTHONUNBUFFERED=1
 
 WORKDIR $HOME/app
 
+COPY --from=builder /usr/local/lib/python3.10/site-packages /usr/local/lib/python3.10/site-packages
 COPY --chown=user . $HOME/app
 
-# Run the application
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,46 +1,62 @@
-from fastapi import FastAPI
-from transformers import AutoModelForCausalLM
-
-# Wrap problematic imports in try-except blocks
-try:
-    from peft import PeftModel, PeftConfig
-except ImportError as e:
-    print(f"Error importing from peft: {e}")
-    raise
-
-try:
-    from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
-except ImportError as e:
-    print(f"Error importing from mistral_common: {e}")
-    raise
+import logging
+from fastapi import FastAPI, HTTPException
+from transformers import AutoModelForCausalLM, pipeline
+from peft import PeftModel, PeftConfig
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load PEFT model configuration and base model
-try:
-    config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
-    base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
-    model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
-
-    # Load recommended tokenizer
-    tokenizer = MistralTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+# Global variables for model, tokenizer, and pipeline
+model = None
+tokenizer = None
+pipe = None
 
-    # Create the pipeline
-    from transformers import pipeline
-    pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
-except Exception as e:
-    print(f"Error loading model or creating pipeline: {e}")
-    raise
+@app.on_event("startup")
+async def load_model():
+    global model, tokenizer, pipe
+    
+    try:
+        logger.info("Loading PEFT configuration...")
+        config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
+        
+        logger.info("Loading base model...")
+        base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+        
+        logger.info("Loading PEFT model...")
+        model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
+
+        logger.info("Loading tokenizer...")
+        tokenizer = MistralTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+
+        logger.info("Creating pipeline...")
+        pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+        
+        logger.info("Model, tokenizer, and pipeline loaded successfully.")
+    except Exception as e:
+        logger.error(f"Error loading model or creating pipeline: {e}")
+        raise
 
 @app.get("/")
 def home():
     return {"message": "Hello World"}
 
 @app.get("/generate")
-def generate(text: str):
+async def generate(text: str):
+    if not pipe:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
     try:
-        output = pipe(text)
+        output = pipe(text, max_length=100, num_return_sequences=1)
         return {"output": output[0]['generated_text']}
     except Exception as e:
-        return {"error": str(e)}
+        logger.error(f"Error during text generation: {e}")
+        raise HTTPException(status_code=500, detail=f"Error during text generation: {str(e)}")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -2,9 +2,9 @@ fastapi==0.103.0
 requests==2.27.*
 uvicorn[standard]==0.17.*
 torch>=1.13.0
-transformers>=4.34.0,<4.35.0
+transformers>=4.34.0,<5.0.0
 numpy<2
-peft>=0.6.0,<0.7.0
+peft>=0.7.0
 accelerate>=0.24.1,<0.25.0
 huggingface_hub>=0.16.4,<0.18.0
 tokenizers>=0.14.0,<0.15.0