Upload 3 files

Dockerfile

@@ -0,0 +1,21 @@
+# Use an official Python base image
+FROM python:3.11-slim
+
+# Set the working directory inside the container
+WORKDIR /code
+
+# Copy the requirements file and install dependencies
+# We use --index-url for the CPU version of torch to keep it smaller
+RUN echo "torch --index-url https://download.pytorch.org/whl/cpu" > /code/requirements.txt
+COPY ./requirements.txt /code/temp_requirements.txt
+RUN cat /code/temp_requirements.txt >> /code/requirements.txt && rm /code/temp_requirements.txt
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# Copy the rest of the application code
+COPY ./app.py /code/app.py
+
+# Expose the port the app runs on (FastAPI default is 8000)
+EXPOSE 8000
+
+# Command to run the application using uvicorn
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

app.py

@@ -0,0 +1,121 @@
+# app.py
+
+import logging
+import time
+from contextlib import asynccontextmanager
+import torch
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# --- Configuration ---
+# Your Hugging Face model repository
+HF_REPO_ID = "rxmha125/RxCodexV1-mini"
+# Use GPU if available (CUDA), otherwise fallback to CPU
+MODEL_LOAD_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# --- Logging Setup ---
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# --- Global variables to hold the model and tokenizer ---
+# These will be loaded during the application's startup.
+model = None
+tokenizer = None
+
+# --- Application Lifespan (Model Loading) ---
+# This special function runs code when the API starts up and shuts down.
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global model, tokenizer
+    logger.info(f"API Startup: Loading model '{HF_REPO_ID}' to device '{MODEL_LOAD_DEVICE}'...")
+    
+    # Load the tokenizer from Hugging Face
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(HF_REPO_ID)
+        # Set a padding token if it's not already set
+        if tokenizer.pad_token is None and tokenizer.eos_token:
+            tokenizer.pad_token = tokenizer.eos_token
+        logger.info("✅ Tokenizer loaded successfully.")
+    except Exception as e:
+        logger.error(f"❌ FATAL: Tokenizer loading failed: {e}")
+
+    # Load the model from Hugging Face
+    try:
+        model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID)
+        model.to(MODEL_LOAD_DEVICE) # Move the model to the correct device (CPU or GPU)
+        model.eval() # Set the model to evaluation mode (important for inference)
+        logger.info("✅ Model loaded successfully.")
+    except Exception as e:
+        logger.error(f"❌ FATAL: Model loading failed: {e}")
+
+    yield # The API is now running
+
+    # --- Code below this line runs on shutdown ---
+    logger.info("API Shutting down.")
+    model = None
+    tokenizer = None
+
+
+# --- Initialize FastAPI ---
+# The 'lifespan' function is linked here to ensure the model loads on startup.
+app = FastAPI(
+    title="Rx Codex V1-mini Simple API",
+    description="A simplified API for text generation without authentication.",
+    lifespan=lifespan
+)
+
+# --- Pydantic Models for API Data Validation ---
+# Defines the structure of the incoming request body
+class GenerationRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 50
+
+# Defines the structure of the outgoing response
+class GenerationResponse(BaseModel):
+    generated_text: str
+
+# --- API Endpoints ---
+@app.get("/")
+def root():
+    """A simple endpoint to check if the API is running."""
+    status = "loaded" if model and tokenizer else "not loaded"
+    return {"message": "Rx Codex API is running", "model_status": status}
+
+@app.post("/generate", response_model=GenerationResponse)
+async def generate_text(request: GenerationRequest):
+    """The main endpoint to generate text from a prompt."""
+    if not model or not tokenizer:
+        raise HTTPException(status_code=503, detail="Model is not ready. Please try again later.")
+    
+    logger.info(f"Received generation request for prompt: '{request.prompt}'")
+    
+    # Prepare the input text for the model
+    inputs = tokenizer(request.prompt, return_tensors="pt", padding=True, truncation=True)
+    input_ids = inputs["input_ids"].to(MODEL_LOAD_DEVICE)
+    
+    # Generate text using the model
+    with torch.no_grad():
+        output_sequences = model.generate(
+            input_ids=input_ids,
+            max_new_tokens=request.max_new_tokens,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+    
+    # Decode the generated tokens back into text
+    generated_text = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
+    
+    # Simple cleanup to remove the original prompt from the output
+    if generated_text.lower().startswith(request.prompt.lower()):
+        generated_text = generated_text[len(request.prompt):].strip()
+
+    logger.info("Generation complete.")
+    return GenerationResponse(generated_text=generated_text)
+
+
+# --- Uvicorn Runner ---
+# This allows you to run the app directly with 'python app.py'
+if __name__ == "__main__":
+    import uvicorn
+    logger.info("Starting API via Uvicorn...")
+    uvicorn.run(app, host="0.0.0.0", port=8000)

requirements.txt

@@ -0,0 +1,10 @@
+# requirements.txt
+
+fastapi
+uvicorn[standard]
+transformers
+sentencepiece
+
+# Install the CPU-only version of PyTorch for smaller size and faster install
+# If you have a powerful NVIDIA GPU, you can remove this line and install the CUDA version
+torch --index-url https://download.pytorch.org/whl/cpu