Initial deployment

Dockerfile

@@ -0,0 +1,31 @@
+# Use NVIDIA CUDA base image for GPU support
+# If you don't have a GPU, use python:3.10-slim
+FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
+
+# Set environment variables
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    python3.10 \
+    python3-pip \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Install Python dependencies
+COPY requirements.txt .
+RUN pip3 install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu118
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY . .
+
+# Expose port (FastAPI default)
+EXPOSE 8000
+
+# Run the application
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

main.py

@@ -0,0 +1,128 @@
+import torch
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline ,BitsAndBytesConfig
+from peft import PeftModel
+from sentence_transformers import SentenceTransformer
+from typing import List, Optional
+import time
+import os
+app = FastAPI(
+    title="Model Deployment API",
+    description="API for contract LoRA generation and text embeddings",
+    version="1.0.0"
+)
+# --- Configuration ---
+LORA_MODEL_ID = "shibinsha02/contract-lora"
+BASE_MODEL_ID = "StevenChen16/llama3-8b-Lawyer"
+EMBEDDING_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
+
+# Global variables for models
+generation_pipeline = None
+embedding_model = SentenceTransformer(EMBEDDING_MODEL_ID, device="cpu")
+# --- Models ---
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_new_tokens: Optional[int] = 128
+    temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 0.9
+
+class GenerateResponse(BaseModel):
+    generated_text: str
+    generation_time: float
+
+class EmbeddingRequest(BaseModel):
+    text: str
+
+class EmbeddingResponse(BaseModel):
+    embedding: List[float]
+    model: str
+
+# --- Startup Event ---
+@app.on_event("startup")
+async def load_models():
+    global generation_pipeline, embedding_model
+    
+    print("Loading embedding model...")
+    embedding_model = SentenceTransformer(EMBEDDING_MODEL_ID)
+    
+    print("Loading generation model (this might take a while)...")
+    # Setting up device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+    
+    # Load tokenizer and base model
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+    
+    # Load with 4-bit quantization if possible for Llama 3 on typical GPUs
+    # Otherwise fallback to float16 or float32
+    try:
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16
+        )
+
+        base_model = AutoModelForCausalLM.from_pretrained(
+            BASE_MODEL_ID,
+            quantization_config=bnb_config,
+            device_map="auto"
+        )
+        # Load LoRA adapter
+        model = PeftModel.from_pretrained(base_model, LORA_MODEL_ID)
+        
+        generation_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device_map="auto" if device == "cuda" else None
+        )
+    except Exception as e:
+        print(f"Error loading generation model: {e}")
+        # Placeholder/Mock for local testing if hardware is insufficient
+        generation_pipeline = None
+
+@app.get("/health")
+async def health_check():
+    return {
+        "status": "healthy",
+        "embeddings_loaded": embedding_model is not None,
+        "generation_loaded": generation_pipeline is not None
+    }
+
+@app.post("/embeddings", response_model=EmbeddingResponse)
+async def get_embeddings(request: EmbeddingRequest):
+    if embedding_model is None:
+        raise HTTPException(status_code=503, detail="Embedding model not loaded")
+    
+    embedding = embedding_model.encode(request.text).tolist()
+    return EmbeddingResponse(
+        embedding=embedding,
+        model=EMBEDDING_MODEL_ID
+    )
+
+@app.post("/generate", response_model=GenerateResponse)
+async def generate_text(request: GenerateRequest):
+    if generation_pipeline is None:
+        raise HTTPException(status_code=503, detail="Generation model not loaded or hardware insufficient")
+    
+    start_time = time.time()
+    
+    outputs = generation_pipeline(
+        request.prompt,
+        max_new_tokens=request.max_new_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        do_sample=True if request.temperature > 0 else False
+    )
+    
+    generated_text = outputs[0]["generated_text"]
+    end_time = time.time()
+    
+    return GenerateResponse(
+        generated_text=generated_text,
+        generation_time=round(end_time - start_time, 2)
+    )
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -0,0 +1,12 @@
+fastapi
+uvicorn
+transformers
+peft
+sentence-transformers
+torch
+pydantic
+accelerate
+bitsandbytes
+python-multipart
+python-dotenv
+httpx