Upload 8 files

Dockerfile

@@ -1,35 +1,25 @@
-# Dockerfile
+# Dockerfile - Ollama style with llama.cpp
 FROM python:3.11-slim
 
-# Set environment variables for Hugging Face cache optimization
 ENV PYTHONUNBUFFERED=1 \
-    PYTHONDONTWRITEBYTECODE=1 \
-    HF_HOME=/tmp/.huggingface \
-    TRANSFORMERS_CACHE=/tmp/.cache/huggingface \
-    HF_HUB_CACHE=/tmp/.cache/huggingface/hub \
-    OMP_NUM_THREADS=4 \
-    MKL_NUM_THREADS=4
+    CMAKE_ARGS="-DLLAMA_AVX2=ON" \
+    FORCE_CMAKE=1
 
-# Install minimal system dependencies
+# System deps for llama.cpp compilation
 RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    cmake \
     git \
     && rm -rf /var/lib/apt/lists/*
 
-# Set working directory
 WORKDIR /app
 
-# Copy requirements first for layer caching
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
+# Install llama-cpp-python (compiles with CPU optimizations)
+RUN pip install --no-cache-dir llama-cpp-python==0.3.2
 
-# Copy application code
+# Copy app
 COPY app/ ./app/
 
-# Create cache directories
-RUN mkdir -p /tmp/.cache/huggingface
-
-# Expose Hugging Face Spaces default port
 EXPOSE 7860
 
-# Run the application
-CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app/model.py

@@ -1,184 +1,196 @@
-# app/model.py
+# app/model.py - llama.cpp optimized version
 """
-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.
+CPU-optimized model loading using llama-cpp-python.
+2-4x faster than transformers on CPU.
 """
 
 import gc
 import os
 from typing import Generator, Optional
+from pathlib import Path
 
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+# Try to use llama.cpp, fallback to transformers
+try:
+    from llama_cpp import Llama
+    LLAMA_AVAILABLE = True
+except ImportError:
+    LLAMA_AVAILABLE = False
+    from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Global singleton instances
-_tokenizer: Optional[AutoTokenizer] = None
-_model: Optional[AutoModelForCausalLM] = None
+# Global singleton
+_llama_model = None
+_transformer_model = None
+_tokenizer = None
 
 
-def load_model() -> tuple[AutoTokenizer, AutoModelForCausalLM]:
+def get_model_path() -> str:
     """
-    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)
+    Returns path to GGUF model.
+    If GGUF not available, returns HF model name.
     """
-    global _tokenizer, _model
-    
-    if _tokenizer is not None and _model is not None:
-        return _tokenizer, _model
-    
-    model_name = "Nanbeige/Nanbeige4.1-3B"
-    
-    # Load tokenizer
-    _tokenizer = AutoTokenizer.from_pretrained(
-        model_name,
-        use_fast=False,
-        trust_remote_code=True
-    )
-    
-    # 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
-    
-    # CPU-optimized model loading
-    # IMPORTANT: Use float32, NOT float16 (float16 is extremely slow on CPU)
-    _model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        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
-    )
+    # Pehle check karo agar GGUF downloaded hai
+    gguf_path = "/tmp/models/nanbeige-3b-q4_0.gguf"
+    if os.path.exists(gguf_path):
+        return gguf_path
     
-    # Explicitly set to CPU (redundant but safe)
-    _model = _model.to("cpu")
-    
-    # Evaluation mode for inference
-    _model.eval()
-    
-    # Clear cache to free memory
-    gc.collect()
-    
-    return _tokenizer, _model
+    # Agar nahi hai, toh HF model name return karo
+    return "Nanbeige/Nanbeige4.1-3B"
 
 
-def generate_stream(
-    prompt: str,
-    temperature: float = 0.7,
-    max_tokens: int = 200
-) -> Generator[str, None, None]:
+def load_model():
     """
-    Generate text in streaming fashion.
-    
-    Args:
-        prompt: Input prompt text
-        temperature: Sampling temperature
-        max_tokens: Maximum tokens to generate
-        
-    Yields:
-        Text chunks as they are generated
+    Load model with llama.cpp if available (GGUF),
+    otherwise fallback to optimized transformers.
     """
-    tokenizer, model = load_model()
-    
-    # Tokenize input
-    inputs = tokenizer(
-        prompt,
-        return_tensors="pt",
-        add_special_tokens=False
-    )
-    
-    # Keep on CPU
-    input_ids = inputs.input_ids
-    
-    # Stream generation using TextIteratorStreamer
-    from transformers import TextIteratorStreamer
-    from threading import Thread
-    
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True
-    )
-    
-    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()
-    
-    for text in streamer:
-        if text:
-            yield text
+    global _llama_model, _transformer_model, _tokenizer
+    
+    # Agar already loaded hai
+    if _llama_model or _transformer_model:
+        return
+    
+    model_path = get_model_path()
+    
+    # GGUF format mein hai toh llama.cpp use karo (FAST)
+    if model_path.endswith(".gguf") and LLAMA_AVAILABLE:
+        print("Loading GGUF model with llama.cpp (optimized)...")
+        _llama_model = Llama(
+            model_path=model_path,
+            n_ctx=2048,
+            n_threads=4,  # CPU threads
+            n_batch=512,
+            verbose=False
+        )
+        print("Model loaded with llama.cpp")
     
-    thread.join()
+    # Nahi toh transformers fallback (SLOW but works)
+    else:
+        print("GGUF not available, using transformers (slower)...")
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        
+        model_name = "Nanbeige/Nanbeige4.1-3B"
+        
+        _tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            use_fast=False
+        )
+        
+        if _tokenizer.pad_token is None:
+            _tokenizer.pad_token = _tokenizer.eos_token
+        
+        _transformer_model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+            device_map=None,
+        )
+        _transformer_model = _transformer_model.to("cpu")
+        _transformer_model.eval()
+        
+        # Disable gradients
+        for param in _transformer_model.parameters():
+            param.requires_grad = False
+        
+        print("Model loaded with transformers")
     
-    # Cleanup
     gc.collect()
 
 
-def generate(
-    prompt: str,
-    temperature: float = 0.7,
-    max_tokens: int = 200
-) -> str:
+def generate_stream(prompt: str, temperature: float = 0.7, max_tokens: int = 100):
     """
-    Generate text non-streaming (full response).
-    
-    Args:
-        prompt: Input prompt text
-        temperature: Sampling temperature
-        max_tokens: Maximum tokens to generate
+    Generate with llama.cpp (fast) or transformers (slow).
+    """
+    load_model()
+    
+    # llama.cpp path (FAST - 2-4x speedup)
+    if _llama_model:
+        # llama.cpp native streaming
+        stream = _llama_model(
+            prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=0.95,
+            stream=True,
+            stop=["</s>", "User:", "Human:"]
+        )
+        
+        for output in stream:
+            text = output["choices"][0]["text"]
+            if text:
+                yield text
+    
+    # Transformers fallback (SLOW)
+    else:
+        import torch
+        from threading import Thread
+        from transformers import TextIteratorStreamer
+        
+        inputs = _tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+        input_ids = inputs.input_ids
+        
+        streamer = TextIteratorStreamer(
+            _tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+        
+        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,
+            "use_cache": True,
+        }
+        
+        thread = Thread(target=_transformer_model.generate, kwargs=generation_kwargs)
+        thread.start()
         
-    Returns:
-        Complete generated text
+        for text in streamer:
+            if text:
+                yield text
+        
+        thread.join()
+        
+    gc.collect()
+
+
+def generate(prompt: str, temperature: float = 0.7, max_tokens: int = 100) -> str:
     """
-    tokenizer, model = load_model()
-    
-    # Tokenize input
-    inputs = tokenizer(
-        prompt,
-        return_tensors="pt",
-        add_special_tokens=False
-    )
-    
-    input_ids = inputs.input_ids
+    Non-streaming generation.
+    """
+    load_model()
     
-    # Generate with no_grad for memory efficiency
-    with torch.no_grad():
-        output_ids = model.generate(
-            input_ids,
-            max_new_tokens=max_tokens,
+    if _llama_model:
+        output = _llama_model(
+            prompt,
+            max_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,
+            stop=["</s>", "User:", "Human:"]
         )
+        return output["choices"][0]["text"]
     
-    # Decode only the new tokens
-    new_tokens = output_ids[0][len(input_ids[0]):]
-    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
-    
-    # Cleanup
-    gc.collect()
-    
-    return response
+    else:
+        import torch
+        inputs = _tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+        
+        with torch.no_grad():
+            output_ids = _transformer_model.generate(
+                inputs.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,
+                use_cache=True,
+            )
+        
+        new_tokens = output_ids[0][len(inputs.input_ids[0]):]
+        return _tokenizer.decode(new_tokens, skip_special_tokens=True)

app/ollama_client.py

@@ -0,0 +1,48 @@
+# app/ollama_client.py
+"""
+Use Ollama if available, otherwise fallback.
+"""
+
+import requests
+import json
+from typing import Generator
+
+
+OLLAMA_URL = "http://localhost:11434"
+
+
+def is_ollama_available() -> bool:
+    try:
+        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=2)
+        return r.status_code == 200
+    except:
+        return False
+
+
+def generate_with_ollama(prompt: str, model: str = "nanbeige", temperature: float = 0.7, max_tokens: int = 100):
+    """
+    Generate using Ollama API (if running).
+    """
+    if not is_ollama_available():
+        raise Exception("Ollama not available")
+    
+    response = requests.post(
+        f"{OLLAMA_URL}/api/generate",
+        json={
+            "model": model,
+            "prompt": prompt,
+            "stream": True,
+            "options": {
+                "temperature": temperature,
+                "num_predict": max_tokens,
+                "top_p": 0.95,
+            }
+        },
+        stream=True
+    )
+    
+    for line in response.iter_lines():
+        if line:
+            data = json.loads(line)
+            if "response" in data:
+                yield data["response"]