Updated app to use ctransformers and gemma without token

app.py

@@ -1,60 +1,86 @@
 import json
 import numpy as np
 from sentence_transformers import SentenceTransformer
-from transformers import pipeline, TextIteratorStreamer
-from threading import Thread
 import torch
 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
+import logging
+import os
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# Load precomputed CV embeddings
-with open("cv_embeddings.json", "r", encoding="utf-8") as f:
-    cv_data = json.load(f)
-    cv_chunks = [item["chunk"] for item in cv_data]
-    cv_embeddings = np.array([item["embedding"] for item in cv_data])
+# 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.")
+    raise ValueError("HF_TOKEN not set")
+login(token=hf_token)  # Set token for huggingface_hub
+
+try:
+    # Load precomputed CV embeddings
+    logger.info("Loading CV embeddings from cv_embeddings.json")
+    with open("cv_embeddings.json", "r", encoding="utf-8") as f:
+        cv_data = json.load(f)
+        cv_chunks = [item["chunk"] for item in cv_data]
+        cv_embeddings = np.array([item["embedding"] for item in cv_data])
+    cv_embeddings_tensor = torch.tensor(cv_embeddings)
+    logger.info("CV embeddings loaded successfully")
 
-cv_embeddings_tensor = torch.tensor(cv_embeddings)
+    # Load embedding model
+    logger.info("Loading SentenceTransformer model")
+    embedder = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
+    logger.info("SentenceTransformer model loaded")
 
-embedder = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
+    # 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",
+    )
+    logger.info("Gemma 3 model loaded")
 
-generator = pipeline(
-    "text-generation",
-    model="distilgpt2",
-    device=-1,
-)
+except Exception as e:
+    logger.error(f"Startup error: {str(e)}", exc_info=True)
+    raise
 
 def retrieve_context(query, top_k=3):
-    query_embedding = embedder.encode(query, convert_to_tensor=True).unsqueeze(0)
-    similarities = F.cosine_similarity(query_embedding, cv_embeddings_tensor, dim=1)
-    top_k = min(top_k, len(similarities))
-    top_indices = torch.topk(similarities, k=top_k).indices.cpu().numpy()
-    return "\n".join([cv_chunks[i] for i in top_indices])
+    try:
+        query_embedding = embedder.encode(query, convert_to_tensor=True).unsqueeze(0)
+        similarities = F.cosine_similarity(query_embedding, cv_embeddings_tensor, dim=1)
+        top_k = min(top_k, len(similarities))
+        top_indices = torch.topk(similarities, k=top_k).indices.cpu().numpy()
+        return "\n".join([cv_chunks[i] for i in top_indices])
+    except Exception as e:
+        logger.error(f"Error in retrieve_context: {str(e)}")
+        raise
 
 def stream_response(query):
-    context = retrieve_context(query)
-    prompt = (
-        f"I am Tim Luka Horstmann, a German Computer Scientist. Based on my CV:\n{context}\n\n"
-        f"Question: {query}\nAnswer:"
-    )
-    
-    streamer = TextIteratorStreamer(generator.tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {
-        "text_inputs": prompt,
-        "max_new_tokens": 200,
-        "do_sample": False,
-        "streamer": streamer,
-    }
-    
-    thread = Thread(target=generator, kwargs=generation_kwargs)
-    thread.start()
-    
-    for token in streamer:
-        yield f"data: {token}\n\n"
-    yield "data: [DONE]\n\n"
+    try:
+        logger.info(f"Processing query: {query}")
+        context = retrieve_context(query)
+        prompt = (
+            f"I am Tim Luka Horstmann, a German Computer Scientist. Based on my CV:\n{context}\n\n"
+            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"
+        yield "data: [DONE]\n\n"
+    except Exception as e:
+        logger.error(f"Error in stream_response: {str(e)}")
+        yield f"data: Error: {str(e)}\n\n"
+        yield "data: [DONE]\n\n"
 
 class QueryRequest(BaseModel):
     data: list
@@ -64,4 +90,8 @@ async def predict(request: QueryRequest):
     if not request.data or not isinstance(request.data, list) or len(request.data) < 1:
         raise HTTPException(status_code=400, detail="Invalid input: 'data' must be a non-empty list")
     query = request.data[0]
-    return StreamingResponse(stream_response(query), media_type="text/event-stream")
+    return StreamingResponse(stream_response(query), media_type="text/event-stream")
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}