Upload folder using huggingface_hub

Dockerfile

@@ -2,14 +2,25 @@ FROM python:3.11-slim
 
 WORKDIR /app
 
+# Set environment variables for CPU optimization
+ENV OMP_NUM_THREADS=2
+ENV MKL_NUM_THREADS=2
+ENV TOKENIZERS_PARALLELISM=true
+ENV TRANSFORMERS_OFFLINE=0
+
 # Install system dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \
     build-essential \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy requirements and install Python dependencies
+# Install PyTorch CPU version first
+RUN pip install --no-cache-dir torch==2.4.1+cpu --extra-index-url https://download.pytorch.org/whl/cpu
+
+# Copy and install other requirements
 COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
+RUN pip install --no-cache-dir fastapi==0.115.0 uvicorn[standard]==0.30.6 \
+    transformers==4.45.0 pydantic==2.9.2 huggingface-hub==0.25.1 \
+    optimum==1.23.0 onnxruntime==1.19.0
 
 # Copy application code
 COPY . .
@@ -18,8 +29,12 @@ COPY . .
 RUN useradd -m -u 1000 user
 USER user
 
+# Pre-download models during build for faster startup
+RUN python -c "from transformers import pipeline; pipeline('text-classification', model='distilbert-base-uncased-finetuned-sst-2-english')" || true
+RUN python -c "from transformers import pipeline; pipeline('text-generation', model='distilgpt2')" || true
+
 # Expose port
 EXPOSE 7860
 
-# Run the application
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+# Run with optimized settings
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]

README.md

@@ -5,12 +5,16 @@ colorFrom: blue
 colorTo: purple
 sdk: docker
 app_port: 7860
+suggested_hardware: cpu-basic
 pinned: false
 ---
 
 # Docker Model Runner
 
-A Docker-based Hugging Face Space with named API endpoints for model inference.
+A CPU-optimized Docker Space with named API endpoints for model inference.
+
+## Hardware
+- **CPU Basic**: 2 vCPU · 16 GB RAM
 
 ## Endpoints
 
@@ -19,27 +23,23 @@ A Docker-based Hugging Face Space with named API endpoints for model inference.
 | `/` | GET | Welcome message |
 | `/health` | GET | Health check |
 | `/info` | GET | Model information |
-| `/predict` | POST | Run model prediction |
+| `/predict` | POST | Text classification |
 | `/generate` | POST | Text generation |
-| `/embed` | POST | Get text embeddings |
+| `/embed` | POST | Text embeddings |
 
 ## Usage
 
-### Health Check
 ```bash
-curl https://YOUR-SPACE.hf.space/health
-```
+# Health Check
+curl https://likhonsheikhdev-docker-model-runner.hf.space/health
 
-### Prediction
-```bash
-curl -X POST https://YOUR-SPACE.hf.space/predict \
+# Prediction
+curl -X POST https://likhonsheikhdev-docker-model-runner.hf.space/predict \
   -H "Content-Type: application/json" \
-  -d '{"text": "Hello, world!"}'
-```
+  -d '{"text": "I love this product!"}'
 
-### Text Generation
-```bash
-curl -X POST https://YOUR-SPACE.hf.space/generate \
+# Text Generation
+curl -X POST https://likhonsheikhdev-docker-model-runner.hf.space/generate \
   -H "Content-Type: application/json" \
   -d '{"prompt": "Once upon a time", "max_length": 50}'
 ```

main.py

@@ -1,5 +1,6 @@
 """
-Docker Model Runner - FastAPI application with named endpoints
+Docker Model Runner - CPU-Optimized FastAPI application
+Optimized for: 2 vCPU, 16GB RAM
 """
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
@@ -8,45 +9,61 @@ import torch
 from transformers import pipeline, AutoTokenizer, AutoModel
 import os
 from datetime import datetime
+from contextlib import asynccontextmanager
 
-app = FastAPI(
-    title="Docker Model Runner",
-    description="HuggingFace Space with named endpoints for model inference",
-    version="1.0.0"
-)
+# CPU-optimized lightweight models
+MODEL_NAME = os.getenv("MODEL_NAME", "distilbert-base-uncased-finetuned-sst-2-english")
+GENERATOR_MODEL = os.getenv("GENERATOR_MODEL", "distilgpt2")
+EMBED_MODEL = os.getenv("EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+
+# Set CPU threading
+torch.set_num_threads(2)
 
-# Model configurations
-MODEL_NAME = os.getenv("MODEL_NAME", "distilbert-base-uncased")
-GENERATOR_MODEL = os.getenv("GENERATOR_MODEL", "gpt2")
+# Global model cache
+models = {}
 
-# Lazy-loaded pipelines
-_classifier = None
-_generator = None
-_embedder = None
 
+def load_models():
+    """Pre-load models for faster inference"""
+    global models
+    print("Loading models for CPU inference...")
 
-def get_classifier():
-    global _classifier
-    if _classifier is None:
-        _classifier = pipeline("text-classification", model=MODEL_NAME)
-    return _classifier
+    # Use smaller, faster models optimized for CPU
+    models["classifier"] = pipeline(
+        "text-classification",
+        model=MODEL_NAME,
+        device=-1,  # CPU
+        torch_dtype=torch.float32
+    )
+
+    models["generator"] = pipeline(
+        "text-generation",
+        model=GENERATOR_MODEL,
+        device=-1,
+        torch_dtype=torch.float32
+    )
 
+    # Lightweight embedding model
+    models["tokenizer"] = AutoTokenizer.from_pretrained(EMBED_MODEL)
+    models["embedder"] = AutoModel.from_pretrained(EMBED_MODEL)
+    models["embedder"].eval()
 
-def get_generator():
-    global _generator
-    if _generator is None:
-        _generator = pipeline("text-generation", model=GENERATOR_MODEL)
-    return _generator
+    print("✅ All models loaded successfully!")
 
 
-def get_embedder():
-    global _embedder
-    if _embedder is None:
-        _embedder = {
-            "tokenizer": AutoTokenizer.from_pretrained(MODEL_NAME),
-            "model": AutoModel.from_pretrained(MODEL_NAME)
-        }
-    return _embedder
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    load_models()
+    yield
+    models.clear()
+
+
+app = FastAPI(
+    title="Docker Model Runner",
+    description="CPU-Optimized HuggingFace Space with named endpoints",
+    version="1.0.0",
+    lifespan=lifespan
+)
 
 
 # Request/Response Models
@@ -65,7 +82,7 @@ class GenerateRequest(BaseModel):
     prompt: str
     max_length: Optional[int] = 50
     num_return_sequences: Optional[int] = 1
-    temperature: Optional[float] = 1.0
+    temperature: Optional[float] = 0.7
 
 
 class GenerateResponse(BaseModel):
@@ -88,12 +105,14 @@ class EmbedResponse(BaseModel):
 class HealthResponse(BaseModel):
     status: str
     timestamp: str
-    gpu_available: bool
+    hardware: str
+    models_loaded: bool
 
 
 class InfoResponse(BaseModel):
     name: str
     version: str
+    hardware: str
     models: dict
     endpoints: List[str]
 
@@ -103,7 +122,8 @@ class InfoResponse(BaseModel):
 async def root():
     """Welcome endpoint"""
     return {
-        "message": "Docker Model Runner API",
+        "message": "Docker Model Runner API (CPU Optimized)",
+        "hardware": "CPU Basic: 2 vCPU · 16 GB RAM",
         "docs": "/docs",
         "endpoints": ["/health", "/info", "/predict", "/generate", "/embed"]
     }
@@ -115,7 +135,8 @@ async def health():
     return HealthResponse(
         status="healthy",
         timestamp=datetime.utcnow().isoformat(),
-        gpu_available=torch.cuda.is_available()
+        hardware="CPU Basic: 2 vCPU · 16 GB RAM",
+        models_loaded=len(models) > 0
     )
 
 
@@ -125,10 +146,11 @@ async def info():
     return InfoResponse(
         name="Docker Model Runner",
         version="1.0.0",
+        hardware="CPU Basic: 2 vCPU · 16 GB RAM",
         models={
             "classifier": MODEL_NAME,
             "generator": GENERATOR_MODEL,
-            "embedder": MODEL_NAME
+            "embedder": EMBED_MODEL
         },
         endpoints=["/", "/health", "/info", "/predict", "/generate", "/embed"]
     )
@@ -137,15 +159,14 @@ async def info():
 @app.post("/predict", response_model=PredictResponse)
 async def predict(request: PredictRequest):
     """
-    Run text classification prediction
+    Run text classification (sentiment analysis)
 
     - **text**: Input text to classify
     - **top_k**: Number of top predictions to return
     """
     try:
         start_time = datetime.now()
-        classifier = get_classifier()
-        results = classifier(request.text, top_k=request.top_k)
+        results = models["classifier"](request.text, top_k=request.top_k)
         latency = (datetime.now() - start_time).total_seconds() * 1000
 
         return PredictResponse(
@@ -163,19 +184,18 @@ async def generate(request: GenerateRequest):
     Generate text from a prompt
 
     - **prompt**: Input prompt for generation
-    - **max_length**: Maximum length of generated text
-    - **num_return_sequences**: Number of sequences to generate
-    - **temperature**: Sampling temperature
+    - **max_length**: Maximum length of generated text (default: 50)
+    - **temperature**: Sampling temperature (default: 0.7)
     """
     try:
         start_time = datetime.now()
-        generator = get_generator()
-        results = generator(
+        results = models["generator"](
             request.prompt,
             max_length=request.max_length,
             num_return_sequences=request.num_return_sequences,
             temperature=request.temperature,
-            do_sample=True
+            do_sample=True,
+            pad_token_id=50256  # GPT2 pad token
         )
         latency = (datetime.now() - start_time).total_seconds() * 1000
 
@@ -193,32 +213,36 @@ async def generate(request: GenerateRequest):
 @app.post("/embed", response_model=EmbedResponse)
 async def embed(request: EmbedRequest):
     """
-    Get text embeddings
+    Get text embeddings using MiniLM (384 dimensions)
 
     - **texts**: List of texts to embed
     """
     try:
         start_time = datetime.now()
-        embedder = get_embedder()
 
-        # Tokenize and get embeddings
-        inputs = embedder["tokenizer"](
+        # Tokenize
+        inputs = models["tokenizer"](
             request.texts,
             padding=True,
             truncation=True,
+            max_length=256,
             return_tensors="pt"
         )
 
+        # Get embeddings
         with torch.no_grad():
-            outputs = embedder["model"](**inputs)
-            # Use mean pooling
-            embeddings = outputs.last_hidden_state.mean(dim=1)
+            outputs = models["embedder"](**inputs)
+            # Mean pooling
+            attention_mask = inputs["attention_mask"]
+            token_embeddings = outputs.last_hidden_state
+            input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+            embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
 
         latency = (datetime.now() - start_time).total_seconds() * 1000
 
         return EmbedResponse(
             embeddings=embeddings.tolist(),
-            model=MODEL_NAME,
+            model=EMBED_MODEL,
             dimensions=embeddings.shape[1],
             latency_ms=round(latency, 2)
         )

requirements.txt

@@ -1,6 +1,8 @@
 fastapi==0.115.0
 uvicorn[standard]==0.30.6
 transformers==4.45.0
-torch==2.4.1
+torch==2.4.1+cpu --extra-index-url https://download.pytorch.org/whl/cpu
 pydantic==2.9.2
 huggingface-hub==0.25.1
+optimum==1.23.0
+onnxruntime==1.19.0