push new model

Dockerfile

@@ -38,7 +38,8 @@ ENV HF_HOME=/models/huggingface \
 # Create cache dir and set permissions
 RUN mkdir -p /models/huggingface && chmod -R 777 /models/huggingface
 
-# Models will be downloaded at runtime (Zephyr is public, no token needed)
+# Pre-download the model during build
+RUN python -c "from transformers import pipeline; import torch; pipe = pipeline('text-generation', model='tiiuae/Falcon3-3B-Instruct', torch_dtype=torch.bfloat16, device_map='cpu')" || true
 
 # Copy project files
 COPY . .

app.py

@@ -9,7 +9,7 @@ from typing import Optional, Tuple
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
 from docx import Document as DocxDocument
 from pptx import Presentation
 import logging
@@ -43,20 +43,19 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-MODEL_ID = "HuggingFaceH4/zephyr-7b-beta"
-tokenizer = None
-model = None
+MODEL_ID = "tiiuae/Falcon3-3B-Instruct"
+pipe = None
 ocr_reader = None
 
 @app.on_event("startup")
 async def load_model():
-    """Load the model, tokenizer, and OCR reader on startup"""
-    global tokenizer, model, ocr_reader
+    """Load the model pipeline and OCR reader on startup"""
+    global pipe, ocr_reader
     try:
         logger.info(f"Loading model: {MODEL_ID} ...")
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
+        pipe = pipeline(
+            "text-generation",
+            model=MODEL_ID,
             torch_dtype=torch.bfloat16,
             device_map="auto"
         )
@@ -381,53 +380,33 @@ Produce ONLY valid JSON with these exact fields:
 }}"""
 
     try:
-        messages = [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": user_message}
-        ]
-        
-        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        
-        prompt_tokens = tokenizer.encode(prompt, return_tensors="pt")
-        prompt_token_count = prompt_tokens.shape[1]
-        
-        max_context = 4096
-        max_input_tokens = 3800
-        
-        if prompt_token_count > max_input_tokens:
-            logger.warning(f"Prompt is {prompt_token_count} tokens, truncating to {max_input_tokens}")
-            prompt_tokens = prompt_tokens[:, :max_input_tokens]
-            prompt = tokenizer.decode(prompt_tokens[0], skip_special_tokens=True)
-            prompt_token_count = max_input_tokens
-        
-        max_output_tokens = max_context - prompt_token_count - 50
+        full_prompt = f"{system_message}\n\n{user_message}"
         
-        logger.info(f"Input tokens: {prompt_token_count}, Available output tokens: {max_output_tokens}")
+        logger.info(f"Input prompt length: {len(full_prompt)} characters")
+        logger.info("Starting model generation with pipeline...")
         
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=prompt_token_count).to(model.device)
-        
-        output_limit = min(1500, max_output_tokens)
-        logger.info(f"Setting max_new_tokens to {output_limit}")
+        messages = [
+            {"role": "user", "content": full_prompt}
+        ]
         
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=output_limit,
+        result = pipe(
+            messages,
+            max_new_tokens=1500,
             temperature=0.3,
             do_sample=True,
             top_p=0.95,
-            pad_token_id=tokenizer.eos_token_id,
-            use_cache=True
+            return_full_text=False
         )
-        raw_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        if "<|assistant|>" in raw_output:
-            raw_output = raw_output.split("<|assistant|>")[-1]
+        raw_output = result[0]["generated_text"]
+        logger.info(f"✅ Generated {len(raw_output)} characters of output")
         
         start = raw_output.find('{')
         end = raw_output.rfind('}') + 1
         
         if start == -1 or end == 0:
-            raise ValueError("No JSON object found in model output")
+            logger.warning("No JSON found in output, returning raw output")
+            raise ValueError(f"No JSON object found in model output. Raw output: {raw_output[:500]}")
         
         parsed_json = json.loads(raw_output[start:end])
         return parsed_json
@@ -458,43 +437,19 @@ Full Deck Length: {len(full_text)} characters
 Produce a FINAL comprehensive review with the same JSON structure as before, consolidating all findings."""
     
     try:
-        messages = [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": user_message}
-        ]
-        
-        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        full_prompt = f"{system_message}\n\n{user_message}"
+        messages = [{"role": "user", "content": full_prompt}]
         
-        prompt_tokens = tokenizer.encode(prompt, return_tensors="pt")
-        prompt_token_count = prompt_tokens.shape[1]
-        
-        max_context = 4096
-        max_input_tokens = 3800
-        
-        if prompt_token_count > max_input_tokens:
-            logger.warning(f"Combine prompt is {prompt_token_count} tokens, truncating to {max_input_tokens}")
-            prompt_tokens = prompt_tokens[:, :max_input_tokens]
-            prompt = tokenizer.decode(prompt_tokens[0], skip_special_tokens=True)
-            prompt_token_count = max_input_tokens
-        
-        max_output_tokens = max_context - prompt_token_count - 50
-        
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=prompt_token_count).to(model.device)
-        
-        output_limit = min(1500, max_output_tokens)
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=output_limit,
+        result = pipe(
+            messages,
+            max_new_tokens=1500,
             temperature=0.3,
             do_sample=True,
             top_p=0.95,
-            pad_token_id=tokenizer.eos_token_id,
-            use_cache=True
+            return_full_text=False
         )
-        raw_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        if "<|assistant|>" in raw_output:
-            raw_output = raw_output.split("<|assistant|>")[-1]
+        raw_output = result[0]["generated_text"]
         
         start = raw_output.find('{')
         end = raw_output.rfind('}') + 1
@@ -583,43 +538,19 @@ Return ONLY valid JSON:
 }}"""
 
     try:
-        messages = [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": user_message}
-        ]
-        
-        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        
-        prompt_tokens = tokenizer.encode(prompt, return_tensors="pt")
-        prompt_token_count = prompt_tokens.shape[1]
-        
-        max_context = 4096
-        max_input_tokens = 3800
-        
-        if prompt_token_count > max_input_tokens:
-            logger.warning(f"Improvement prompt is {prompt_token_count} tokens, truncating to {max_input_tokens}")
-            prompt_tokens = prompt_tokens[:, :max_input_tokens]
-            prompt = tokenizer.decode(prompt_tokens[0], skip_special_tokens=True)
-            prompt_token_count = max_input_tokens
-        
-        max_output_tokens = max_context - prompt_token_count - 50
-        
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=prompt_token_count).to(model.device)
+        full_prompt = f"{system_message}\n\n{user_message}"
+        messages = [{"role": "user", "content": full_prompt}]
         
-        output_limit = min(1000, max_output_tokens)
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=output_limit,
+        result = pipe(
+            messages,
+            max_new_tokens=1000,
             temperature=0.4,
             do_sample=True,
             top_p=0.95,
-            pad_token_id=tokenizer.eos_token_id,
-            use_cache=True
+            return_full_text=False
         )
-        raw_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        if "<|assistant|>" in raw_output:
-            raw_output = raw_output.split("<|assistant|>")[-1]
+        raw_output = result[0]["generated_text"]
         
         start = raw_output.find('{')
         end = raw_output.rfind('}') + 1
@@ -656,7 +587,7 @@ async def health():
     """Health check endpoint"""
     return {
         "status": "healthy",
-        "model_loaded": model is not None and tokenizer is not None
+        "model_loaded": pipe is not None
     }
 
 @app.post("/review")
@@ -666,7 +597,7 @@ async def review_deck(file: UploadFile = File(...)):
     
     Supported formats: PDF, DOCX, PPT, PPTX
     """
-    if model is None or tokenizer is None:
+    if pipe is None:
         raise HTTPException(status_code=503, detail="Model not loaded yet. Please wait for startup to complete.")
     
     file_extension = Path(file.filename).suffix.lower()