updated to llama

Dockerfile

@@ -4,14 +4,16 @@ FROM python:3.10
 # Set working directory
 WORKDIR /app
 
-# Install system dependencies for ctransformers and runtime
+# Install system dependencies for llama_cpp
 RUN apt-get update && apt-get install -y \
     gcc \
     g++ \
     libffi-dev \
+    libgcc-s1 \
+    libstdc++6 \
     && rm -rf /var/lib/apt/lists/*
 
-# Set environment variables for cache and token
+# Set environment variables for cache
 ENV TRANSFORMERS_CACHE=/app/cache
 ENV HF_HOME=/app/cache
 

app.py

@@ -6,8 +6,8 @@ import torch.nn.functional as F
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-from ctransformers import AutoModelForCausalLM
-from huggingface_hub import login
+from llama_cpp import Llama
+from huggingface_hub import login, hf_hub_download
 import logging
 import os
 
@@ -17,7 +17,7 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# Authenticate with Hugging Fac
+# Authenticate with Hugging Face
 hf_token = os.getenv("HF_TOKEN")
 if not hf_token:
     logger.error("HF_TOKEN environment variable not set. Required for gated models.")
@@ -39,15 +39,20 @@ try:
     embedder = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
     logger.info("SentenceTransformer model loaded")
 
-    # Load Gemma 3 model with ctransformers
-    logger.info("Loading Gemma 3 model")
-    generator = AutoModelForCausalLM.from_pretrained(
-        "google/gemma-3-12b-it-qat-q4_0-gguf",
-        local_files_only=False,
-        model_type="gemma",
-        model_file="gemma-3-12b-it-q4_0.gguf",
+    # Load Gemma 3 1B model with llama_cpp
+    logger.info("Loading Gemma 3 1B model")
+    model_path = hf_hub_download(
+        repo_id="google/gemma-3-1b-it-qat-q4_0-gguf",
+        filename="gemma-3-1b-it-q4_0.gguf",
+        local_dir="/app/cache" if os.getenv("HF_HOME") else None,  # Use cache dir in Docker
+        token=hf_token,
     )
-    logger.info("Gemma 3 model loaded")
+    generator = Llama(
+        model_path=model_path,
+        n_ctx=2048,  # Context length
+        n_threads=4,  # Adjust based on CPU cores
+    )
+    logger.info("Gemma 3 1B model loaded")
 
 except Exception as e:
     logger.error(f"Startup error: {str(e)}", exc_info=True)
@@ -73,9 +78,14 @@ def stream_response(query):
             f"Question: {query}\nAnswer:"
         )
         
-        # Stream response with ctransformers
-        for token in generator(prompt, max_new_tokens=512, stream=True):
-            yield f"data: {token}\n\n"
+        # Stream response with llama_cpp
+        for chunk in generator(
+            prompt,
+            max_tokens=512,
+            stream=True,
+            stop=["[DONE]"],
+        ):
+            yield f"data: {chunk['choices'][0]['text']}\n\n"
         yield "data: [DONE]\n\n"
     except Exception as e:
         logger.error(f"Error in stream_response: {str(e)}")

requirements.txt

@@ -1,7 +1,7 @@
-fastapi
-uvicorn
-sentence-transformers
-torch
-numpy
-ctransformers
-huggingface_hub
+fastapi==0.115.0
+uvicorn==0.31.0
+sentence-transformers==3.1.1
+torch==2.4.1
+numpy==1.26.4
+llama-cpp-python==0.3.1
+huggingface_hub==0.30.1