Upload 7 files

Dockerfile

@@ -6,7 +6,9 @@ ENV PYTHONUNBUFFERED=1 \
     PYTHONDONTWRITEBYTECODE=1 \
     HF_HOME=/tmp/.huggingface \
     TRANSFORMERS_CACHE=/tmp/.cache/huggingface \
-    HF_HUB_CACHE=/tmp/.cache/huggingface/hub
+    HF_HUB_CACHE=/tmp/.cache/huggingface/hub \
+    OMP_NUM_THREADS=4 \
+    MKL_NUM_THREADS=4
 
 # Install minimal system dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \

app/main.py

@@ -1,7 +1,7 @@
 # app/main.py
 """
 FastAPI application for serving Nanbeige4.1-3B model.
-Optimized for Hugging Face Spaces (CPU, Docker).
+CPU-ONLY optimized for Hugging Face Spaces (Docker).
 """
 
 import asyncio
@@ -22,17 +22,17 @@ async def lifespan(app: FastAPI):
     Loads model on startup to ensure it's ready for requests.
     """
     # Startup: Load model
-    print("Loading model...")
+    print("Loading model on CPU...")
     load_model()
-    print("Model loaded successfully")
+    print("Model loaded successfully on CPU")
     yield
-    # Shutdown: Cleanup (if needed)
+    # Shutdown: Cleanup
     print("Shutting down...")
 
 
 app = FastAPI(
-    title="Nanbeige4.1-3B API",
-    description="FastAPI wrapper for Nanbeige4.1-3B with streaming support",
+    title="Nanbeige4.1-3B API (CPU)",
+    description="FastAPI wrapper for Nanbeige4.1-3B - CPU Optimized",
     version="1.0.0",
     lifespan=lifespan
 )
@@ -41,7 +41,12 @@ app = FastAPI(
 @app.get("/")
 async def health_check():
     """Health check endpoint."""
-    return {"status": "ok", "model": "Nanbeige4.1-3B"}
+    return {
+        "status": "ok", 
+        "model": "Nanbeige4.1-3B",
+        "device": "cpu",
+        "mode": "float32"
+    }
 
 
 @app.post("/generate")
@@ -56,14 +61,21 @@ async def generate_text(request: GenerationRequest):
     if request.stream:
         # Streaming response
         async def stream_generator():
-            # Run sync generator in thread pool to not block
+            # Run sync generator in thread pool to not block event loop
             loop = asyncio.get_event_loop()
-            sync_gen = generate_stream(
-                final_prompt,
-                temperature=request.temperature,
-                max_tokens=request.max_tokens
-            )
             
+            # Use run_in_executor for CPU-bound operations
+            def sync_generator():
+                return generate_stream(
+                    final_prompt,
+                    temperature=request.temperature,
+                    max_tokens=request.max_tokens
+                )
+            
+            # Get the generator
+            sync_gen = await loop.run_in_executor(None, sync_generator)
+            
+            # Iterate through chunks
             for chunk in sync_gen:
                 if chunk:
                     # SSE format
@@ -80,10 +92,14 @@ async def generate_text(request: GenerationRequest):
             }
         )
     else:
-        # Non-streaming response
-        result = generate(
-            final_prompt,
-            temperature=request.temperature,
-            max_tokens=request.max_tokens
+        # Non-streaming response - run in executor to not block
+        loop = asyncio.get_event_loop()
+        result = await loop.run_in_executor(
+            None,
+            lambda: generate(
+                final_prompt,
+                temperature=request.temperature,
+                max_tokens=request.max_tokens
+            )
         )
         return GenerationResponse(text=result)

app/model.py

@@ -1,6 +1,7 @@
 # app/model.py
 """
 Model loading and inference utilities for Nanbeige/Nanbeige4.1-3B.
+CPU-optimized implementation - NO GPU/CUDA code.
 Implements singleton pattern to ensure model loads only once.
 """
 
@@ -9,35 +10,24 @@ import os
 from typing import Generator, Optional
 
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Global singleton instances
 _tokenizer: Optional[AutoTokenizer] = None
 _model: Optional[AutoModelForCausalLM] = None
 
 
-def get_quantization_config() -> Optional[BitsAndBytesConfig]:
-    """
-    Configure 4-bit quantization for CPU memory efficiency.
-    Returns None if bitsandbytes is not available or on CPU.
-    """
-    try:
-        # 4-bit quantization config for minimal memory footprint
-        return BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.float16,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_use_double_quant=True,
-        )
-    except Exception:
-        return None
-
-
 def load_model() -> tuple[AutoTokenizer, AutoModelForCausalLM]:
     """
     Load tokenizer and model with singleton pattern.
     Loads only on first call, returns cached instances thereafter.
     
+    CPU Optimization Notes:
+    - Use torch.float32 (float16 is 7x slower on CPU)
+    - low_cpu_mem_usage=True prevents memory spikes
+    - No device_map (CPU pe auto mat use karna)
+    - trust_remote_code=True required for Nanbeige models
+    
     Returns:
         Tuple of (tokenizer, model)
     """
@@ -55,32 +45,29 @@ def load_model() -> tuple[AutoTokenizer, AutoModelForCausalLM]:
         trust_remote_code=True
     )
     
-    # Configure model loading for CPU
-    # Use torch.float16 for memory efficiency on CPU
-    model_kwargs = {
-        "torch_dtype": torch.float16,
-        "trust_remote_code": True,
-        "low_cpu_mem_usage": True,
-    }
-    
-    # Try to use quantization if available, otherwise use standard loading
-    quant_config = get_quantization_config()
-    if quant_config is not None:
-        model_kwargs["quantization_config"] = quant_config
+    # Set pad token if not present
+    if _tokenizer.pad_token is None:
+        _tokenizer.pad_token = _tokenizer.eos_token
+        _tokenizer.pad_token_id = _tokenizer.eos_token_id
     
-    # Load model
+    # CPU-optimized model loading
+    # IMPORTANT: Use float32, NOT float16 (float16 is extremely slow on CPU)
     _model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        **model_kwargs
+        torch_dtype=torch.float32,  # CPU pe float32 best hai
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,     # Memory optimization
+        device_map=None,            # CPU pe explicit None rakho
     )
     
-    # Ensure model is in eval mode
+    # Explicitly set to CPU (redundant but safe)
+    _model = _model.to("cpu")
+    
+    # Evaluation mode for inference
     _model.eval()
     
     # Clear cache to free memory
     gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
     
     return _tokenizer, _model
 
@@ -110,21 +97,10 @@ def generate_stream(
         add_special_tokens=False
     )
     
-    # Move to same device as model
-    input_ids = inputs.input_ids.to(model.device)
+    # Keep on CPU
+    input_ids = inputs.input_ids
     
-    # Generation parameters optimized for Nanbeige
-    generation_kwargs = {
-        "input_ids": input_ids,
-        "max_new_tokens": max_tokens,
-        "temperature": temperature,
-        "top_p": 0.95,
-        "do_sample": True,
-        "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
-        "eos_token_id": tokenizer.eos_token_id,
-    }
-    
-    # Stream generation using generate with streamer
+    # Stream generation using TextIteratorStreamer
     from transformers import TextIteratorStreamer
     from threading import Thread
     
@@ -133,16 +109,25 @@ def generate_stream(
         skip_prompt=True,
         skip_special_tokens=True
     )
-    generation_kwargs["streamer"] = streamer
+    
+    generation_kwargs = {
+        "input_ids": input_ids,
+        "max_new_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": 0.95,
+        "do_sample": True,
+        "pad_token_id": tokenizer.pad_token_id,
+        "eos_token_id": tokenizer.eos_token_id,
+        "streamer": streamer,
+    }
     
     # Run generation in separate thread to enable streaming
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     
-    generated_text = ""
     for text in streamer:
-        generated_text += text
-        yield text
+        if text:
+            yield text
     
     thread.join()
     
@@ -175,9 +160,9 @@ def generate(
         add_special_tokens=False
     )
     
-    input_ids = inputs.input_ids.to(model.device)
+    input_ids = inputs.input_ids
     
-    # Generate
+    # Generate with no_grad for memory efficiency
     with torch.no_grad():
         output_ids = model.generate(
             input_ids,
@@ -185,7 +170,7 @@ def generate(
             temperature=temperature,
             top_p=0.95,
             do_sample=True,
-            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
         )
     

app/schemas.py

@@ -23,7 +23,7 @@ class GenerationRequest(BaseModel):
     max_tokens: int = Field(
         default=200,
         ge=1,
-        le=2048,
+        le=512,
         description="Maximum tokens to generate"
     )
     stream: bool = Field(

requirements.txt

@@ -6,6 +6,5 @@ transformers==4.46.0
 torch==2.5.0
 accelerate==1.0.0
 sentencepiece==0.2.0
-bitsandbytes==0.44.0
 huggingface-hub==0.26.0
 python-multipart==0.0.12

run.sh

@@ -1,6 +1,11 @@
 #!/bin/bash
 # run.sh
 # Production startup script for uvicorn server
+# Optimized for CPU-only Hugging Face Spaces
+
+export OMP_NUM_THREADS=4
+export MKL_NUM_THREADS=4
+export TRANSFORMERS_CACHE=/tmp/.cache/huggingface
 
 exec uvicorn app.main:app \
     --host 0.0.0.0 \