Upload folder using huggingface_hub

README.md

@@ -11,35 +11,83 @@ pinned: false
 
 # Docker Model Runner
 
-A CPU-optimized Docker Space with named API endpoints for model inference.
+Anthropic & OpenAI API compatible Docker Space with named endpoints.
 
 ## Hardware
 - **CPU Basic**: 2 vCPU · 16 GB RAM
 
+## API Compatibility
+
+### Anthropic Messages API
+```bash
+curl -X POST https://likhonsheikhdev-docker-model-runner.hf.space/v1/messages \
+  -H "Content-Type: application/json" \
+  -H "x-api-key: your-key" \
+  -d '{
+    "model": "distilgpt2",
+    "max_tokens": 256,
+    "messages": [
+      {"role": "user", "content": "Hello, how are you?"}
+    ]
+  }'
+```
+
+### OpenAI Chat Completions API
+```bash
+curl -X POST https://likhonsheikhdev-docker-model-runner.hf.space/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -H "Authorization: Bearer your-key" \
+  -d '{
+    "model": "distilgpt2",
+    "messages": [
+      {"role": "user", "content": "Hello, how are you?"}
+    ]
+  }'
+```
+
 ## Endpoints
 
 | Endpoint | Method | Description |
 |----------|--------|-------------|
-| `/` | GET | Welcome message |
+| `/v1/messages` | POST | Anthropic Messages API |
+| `/v1/chat/completions` | POST | OpenAI Chat API |
+| `/v1/models` | GET | List available models |
 | `/health` | GET | Health check |
-| `/info` | GET | Model information |
+| `/info` | GET | API information |
 | `/predict` | POST | Text classification |
-| `/generate` | POST | Text generation |
 | `/embed` | POST | Text embeddings |
 
-## Usage
+## Python SDK Usage
 
-```bash
-# Health Check
-curl https://likhonsheikhdev-docker-model-runner.hf.space/health
+### With Anthropic SDK
+```python
+from anthropic import Anthropic
 
-# Prediction
-curl -X POST https://likhonsheikhdev-docker-model-runner.hf.space/predict \
-  -H "Content-Type: application/json" \
-  -d '{"text": "I love this product!"}'
+client = Anthropic(
+    api_key="any-key",
+    base_url="https://likhonsheikhdev-docker-model-runner.hf.space"
+)
 
-# Text Generation
-curl -X POST https://likhonsheikhdev-docker-model-runner.hf.space/generate \
-  -H "Content-Type: application/json" \
-  -d '{"prompt": "Once upon a time", "max_length": 50}'
+message = client.messages.create(
+    model="distilgpt2",
+    max_tokens=256,
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+print(message.content[0].text)
+```
+
+### With OpenAI SDK
+```python
+from openai import OpenAI
+
+client = OpenAI(
+    api_key="any-key",
+    base_url="https://likhonsheikhdev-docker-model-runner.hf.space/v1"
+)
+
+response = client.chat.completions.create(
+    model="distilgpt2",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+print(response.choices[0].message.content)
 ```

main.py

@@ -1,15 +1,18 @@
 """
 Docker Model Runner - CPU-Optimized FastAPI application
+Compatible with Anthropic API format
 Optimized for: 2 vCPU, 16GB RAM
 """
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from typing import Optional, List
+from fastapi import FastAPI, HTTPException, Header
+from pydantic import BaseModel, Field
+from typing import Optional, List, Union, Literal
 import torch
-from transformers import pipeline, AutoTokenizer, AutoModel
+from transformers import pipeline, AutoTokenizer, AutoModel, AutoModelForCausalLM
 import os
 from datetime import datetime
 from contextlib import asynccontextmanager
+import uuid
+import time
 
 # CPU-optimized lightweight models
 MODEL_NAME = os.getenv("MODEL_NAME", "distilbert-base-uncased-finetuned-sst-2-english")
@@ -28,25 +31,27 @@ def load_models():
     global models
     print("Loading models for CPU inference...")
 
-    # Use smaller, faster models optimized for CPU
+    # Classifier
     models["classifier"] = pipeline(
         "text-classification",
         model=MODEL_NAME,
-        device=-1,  # CPU
-        torch_dtype=torch.float32
-    )
-
-    models["generator"] = pipeline(
-        "text-generation",
-        model=GENERATOR_MODEL,
         device=-1,
         torch_dtype=torch.float32
     )
 
-    # Lightweight embedding model
-    models["tokenizer"] = AutoTokenizer.from_pretrained(EMBED_MODEL)
-    models["embedder"] = AutoModel.from_pretrained(EMBED_MODEL)
-    models["embedder"].eval()
+    # Generator with tokenizer for chat
+    models["generator_tokenizer"] = AutoTokenizer.from_pretrained(GENERATOR_MODEL)
+    models["generator_model"] = AutoModelForCausalLM.from_pretrained(GENERATOR_MODEL)
+    models["generator_model"].eval()
+
+    # Set pad token
+    if models["generator_tokenizer"].pad_token is None:
+        models["generator_tokenizer"].pad_token = models["generator_tokenizer"].eos_token
+
+    # Embedding model
+    models["embed_tokenizer"] = AutoTokenizer.from_pretrained(EMBED_MODEL)
+    models["embed_model"] = AutoModel.from_pretrained(EMBED_MODEL)
+    models["embed_model"].eval()
 
     print("✅ All models loaded successfully!")
 
@@ -60,33 +65,91 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(
     title="Docker Model Runner",
-    description="CPU-Optimized HuggingFace Space with named endpoints",
+    description="Anthropic API Compatible - CPU-Optimized HuggingFace Space",
     version="1.0.0",
     lifespan=lifespan
 )
 
 
-# Request/Response Models
-class PredictRequest(BaseModel):
+# ============== Anthropic API Models ==============
+
+class ContentBlock(BaseModel):
+    type: Literal["text"] = "text"
     text: str
-    top_k: Optional[int] = 1
 
 
-class PredictResponse(BaseModel):
-    predictions: List[dict]
+class MessageContent(BaseModel):
+    role: Literal["user", "assistant"]
+    content: Union[str, List[ContentBlock]]
+
+
+class AnthropicRequest(BaseModel):
+    model: str = "distilgpt2"
+    messages: List[MessageContent]
+    max_tokens: int = 1024
+    temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 1.0
+    stop_sequences: Optional[List[str]] = None
+    stream: Optional[bool] = False
+    system: Optional[str] = None
+
+
+class Usage(BaseModel):
+    input_tokens: int
+    output_tokens: int
+
+
+class AnthropicResponse(BaseModel):
+    id: str
+    type: Literal["message"] = "message"
+    role: Literal["assistant"] = "assistant"
+    content: List[ContentBlock]
     model: str
-    latency_ms: float
+    stop_reason: Literal["end_turn", "max_tokens", "stop_sequence"] = "end_turn"
+    stop_sequence: Optional[str] = None
+    usage: Usage
+
+
+# ============== OpenAI Compatible Models ==============
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
 
 
-class GenerateRequest(BaseModel):
-    prompt: str
-    max_length: Optional[int] = 50
-    num_return_sequences: Optional[int] = 1
+class ChatCompletionRequest(BaseModel):
+    model: str = "distilgpt2"
+    messages: List[ChatMessage]
+    max_tokens: Optional[int] = 1024
     temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 1.0
+    stream: Optional[bool] = False
 
 
-class GenerateResponse(BaseModel):
-    generated_text: List[str]
+class ChatChoice(BaseModel):
+    index: int = 0
+    message: ChatMessage
+    finish_reason: str = "stop"
+
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[ChatChoice]
+    usage: dict
+
+
+# ============== Other Request/Response Models ==============
+
+class PredictRequest(BaseModel):
+    text: str
+    top_k: Optional[int] = 1
+
+
+class PredictResponse(BaseModel):
+    predictions: List[dict]
     model: str
     latency_ms: float
 
@@ -109,23 +172,178 @@ class HealthResponse(BaseModel):
     models_loaded: bool
 
 
-class InfoResponse(BaseModel):
-    name: str
-    version: str
-    hardware: str
-    models: dict
-    endpoints: List[str]
+class ModelInfo(BaseModel):
+    id: str
+    object: str = "model"
+    created: int
+    owned_by: str = "local"
+
+
+class ModelsResponse(BaseModel):
+    object: str = "list"
+    data: List[ModelInfo]
+
+
+# ============== Helper Functions ==============
+
+def generate_text(prompt: str, max_tokens: int, temperature: float, top_p: float) -> tuple:
+    """Generate text and return (text, input_tokens, output_tokens)"""
+    tokenizer = models["generator_tokenizer"]
+    model = models["generator_model"]
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+    input_tokens = inputs["input_ids"].shape[1]
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            temperature=temperature if temperature > 0 else 1.0,
+            top_p=top_p,
+            do_sample=temperature > 0,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
+
+    generated_tokens = outputs[0][input_tokens:]
+    output_tokens = len(generated_tokens)
+    generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+
+    return generated_text.strip(), input_tokens, output_tokens
+
+
+def format_messages_to_prompt(messages: List, system: Optional[str] = None) -> str:
+    """Convert chat messages to a single prompt string"""
+    prompt_parts = []
+
+    if system:
+        prompt_parts.append(f"System: {system}\n")
+
+    for msg in messages:
+        role = msg.role if hasattr(msg, 'role') else msg.get('role', 'user')
+        content = msg.content if hasattr(msg, 'content') else msg.get('content', '')
+
+        # Handle content that might be a list of blocks
+        if isinstance(content, list):
+            content = " ".join([block.text if hasattr(block, 'text') else block.get('text', '') for block in content])
+
+        if role == "user":
+            prompt_parts.append(f"Human: {content}\n")
+        elif role == "assistant":
+            prompt_parts.append(f"Assistant: {content}\n")
+
+    prompt_parts.append("Assistant:")
+    return "".join(prompt_parts)
+
+
+# ============== Anthropic API Endpoints ==============
+
+@app.post("/v1/messages", response_model=AnthropicResponse)
+async def create_message(
+    request: AnthropicRequest,
+    x_api_key: Optional[str] = Header(None, alias="x-api-key"),
+    authorization: Optional[str] = Header(None)
+):
+    """
+    Anthropic Messages API compatible endpoint
+
+    POST /v1/messages
+    """
+    try:
+        # Format messages to prompt
+        prompt = format_messages_to_prompt(request.messages, request.system)
+
+        # Generate response
+        generated_text, input_tokens, output_tokens = generate_text(
+            prompt=prompt,
+            max_tokens=request.max_tokens,
+            temperature=request.temperature or 0.7,
+            top_p=request.top_p or 1.0
+        )
+
+        return AnthropicResponse(
+            id=f"msg_{uuid.uuid4().hex[:24]}",
+            content=[ContentBlock(type="text", text=generated_text)],
+            model=GENERATOR_MODEL,
+            stop_reason="end_turn",
+            usage=Usage(input_tokens=input_tokens, output_tokens=output_tokens)
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+# ============== OpenAI Compatible Endpoints ==============
+
+@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+async def chat_completions(
+    request: ChatCompletionRequest,
+    authorization: Optional[str] = Header(None)
+):
+    """
+    OpenAI Chat Completions API compatible endpoint
 
+    POST /v1/chat/completions
+    """
+    try:
+        # Format messages to prompt
+        prompt = format_messages_to_prompt(request.messages)
+
+        # Generate response
+        generated_text, input_tokens, output_tokens = generate_text(
+            prompt=prompt,
+            max_tokens=request.max_tokens or 1024,
+            temperature=request.temperature or 0.7,
+            top_p=request.top_p or 1.0
+        )
+
+        return ChatCompletionResponse(
+            id=f"chatcmpl-{uuid.uuid4().hex[:24]}",
+            created=int(time.time()),
+            model=GENERATOR_MODEL,
+            choices=[
+                ChatChoice(
+                    index=0,
+                    message=ChatMessage(role="assistant", content=generated_text),
+                    finish_reason="stop"
+                )
+            ],
+            usage={
+                "prompt_tokens": input_tokens,
+                "completion_tokens": output_tokens,
+                "total_tokens": input_tokens + output_tokens
+            }
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/v1/models", response_model=ModelsResponse)
+async def list_models():
+    """List available models (OpenAI compatible)"""
+    return ModelsResponse(
+        data=[
+            ModelInfo(id=GENERATOR_MODEL, created=int(time.time())),
+            ModelInfo(id=MODEL_NAME, created=int(time.time())),
+            ModelInfo(id=EMBED_MODEL, created=int(time.time()))
+        ]
+    )
+
+
+# ============== Original Endpoints ==============
 
-# Named Endpoints
 @app.get("/")
 async def root():
     """Welcome endpoint"""
     return {
-        "message": "Docker Model Runner API (CPU Optimized)",
+        "message": "Docker Model Runner API (Anthropic Compatible)",
         "hardware": "CPU Basic: 2 vCPU · 16 GB RAM",
         "docs": "/docs",
-        "endpoints": ["/health", "/info", "/predict", "/generate", "/embed"]
+        "api_endpoints": {
+            "anthropic": "/v1/messages",
+            "openai": "/v1/chat/completions",
+            "models": "/v1/models"
+        },
+        "utility_endpoints": ["/health", "/info", "/predict", "/embed"]
     }
 
 
@@ -140,30 +358,32 @@ async def health():
     )
 
 
-@app.get("/info", response_model=InfoResponse)
+@app.get("/info")
 async def info():
     """Model and API information"""
-    return InfoResponse(
-        name="Docker Model Runner",
-        version="1.0.0",
-        hardware="CPU Basic: 2 vCPU · 16 GB RAM",
-        models={
+    return {
+        "name": "Docker Model Runner",
+        "version": "1.0.0",
+        "api_compatibility": ["anthropic", "openai"],
+        "hardware": "CPU Basic: 2 vCPU · 16 GB RAM",
+        "models": {
+            "chat": GENERATOR_MODEL,
             "classifier": MODEL_NAME,
-            "generator": GENERATOR_MODEL,
             "embedder": EMBED_MODEL
         },
-        endpoints=["/", "/health", "/info", "/predict", "/generate", "/embed"]
-    )
+        "endpoints": {
+            "anthropic_messages": "POST /v1/messages",
+            "openai_chat": "POST /v1/chat/completions",
+            "models": "GET /v1/models",
+            "predict": "POST /predict",
+            "embed": "POST /embed"
+        }
+    }
 
 
 @app.post("/predict", response_model=PredictResponse)
 async def predict(request: PredictRequest):
-    """
-    Run text classification (sentiment analysis)
-
-    - **text**: Input text to classify
-    - **top_k**: Number of top predictions to return
-    """
+    """Text classification (sentiment analysis)"""
     try:
         start_time = datetime.now()
         results = models["classifier"](request.text, top_k=request.top_k)
@@ -178,50 +398,13 @@ async def predict(request: PredictRequest):
         raise HTTPException(status_code=500, detail=str(e))
 
 
-@app.post("/generate", response_model=GenerateResponse)
-async def generate(request: GenerateRequest):
-    """
-    Generate text from a prompt
-
-    - **prompt**: Input prompt for generation
-    - **max_length**: Maximum length of generated text (default: 50)
-    - **temperature**: Sampling temperature (default: 0.7)
-    """
-    try:
-        start_time = datetime.now()
-        results = models["generator"](
-            request.prompt,
-            max_length=request.max_length,
-            num_return_sequences=request.num_return_sequences,
-            temperature=request.temperature,
-            do_sample=True,
-            pad_token_id=50256  # GPT2 pad token
-        )
-        latency = (datetime.now() - start_time).total_seconds() * 1000
-
-        generated_texts = [r["generated_text"] for r in results]
-
-        return GenerateResponse(
-            generated_text=generated_texts,
-            model=GENERATOR_MODEL,
-            latency_ms=round(latency, 2)
-        )
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-
 @app.post("/embed", response_model=EmbedResponse)
 async def embed(request: EmbedRequest):
-    """
-    Get text embeddings using MiniLM (384 dimensions)
-
-    - **texts**: List of texts to embed
-    """
+    """Get text embeddings"""
     try:
         start_time = datetime.now()
 
-        # Tokenize
-        inputs = models["tokenizer"](
+        inputs = models["embed_tokenizer"](
             request.texts,
             padding=True,
             truncation=True,
@@ -229,10 +412,8 @@ async def embed(request: EmbedRequest):
             return_tensors="pt"
         )
 
-        # Get embeddings
         with torch.no_grad():
-            outputs = models["embedder"](**inputs)
-            # Mean pooling
+            outputs = models["embed_model"](**inputs)
             attention_mask = inputs["attention_mask"]
             token_embeddings = outputs.last_hidden_state
             input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()