Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,326 @@
+"""
+Anthropic-Compatible API Endpoint
+Lightweight CPU-based implementation for Hugging Face Spaces
+"""
+
+import os
+import time
+import uuid
+from typing import List, Optional, Union
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI, HTTPException, Header, Request
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+import json
+
+# ============== Configuration ==============
+MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"  # Ultra-lightweight 135M model
+MAX_TOKENS_DEFAULT = 1024
+DEVICE = "cpu"
+
+# Global model and tokenizer
+model = None
+tokenizer = None
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Load model on startup"""
+    global model, tokenizer
+    print(f"Loading model: {MODEL_ID}")
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float32,
+        device_map=DEVICE,
+        low_cpu_mem_usage=True
+    )
+    model.eval()
+    print("Model loaded successfully!")
+
+    yield
+
+    # Cleanup
+    del model, tokenizer
+
+app = FastAPI(
+    title="Anthropic-Compatible API",
+    description="Lightweight CPU-based API with Anthropic Messages API compatibility",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ============== Pydantic Models (Anthropic-Compatible) ==============
+
+class ContentBlock(BaseModel):
+    type: str = "text"
+    text: str
+
+class Message(BaseModel):
+    role: str
+    content: Union[str, List[ContentBlock]]
+
+class MessageRequest(BaseModel):
+    model: str
+    messages: List[Message]
+    max_tokens: int = MAX_TOKENS_DEFAULT
+    temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 0.9
+    top_k: Optional[int] = 50
+    stream: Optional[bool] = False
+    system: Optional[str] = None
+    stop_sequences: Optional[List[str]] = None
+
+class Usage(BaseModel):
+    input_tokens: int
+    output_tokens: int
+
+class MessageResponse(BaseModel):
+    id: str
+    type: str = "message"
+    role: str = "assistant"
+    content: List[ContentBlock]
+    model: str
+    stop_reason: str = "end_turn"
+    stop_sequence: Optional[str] = None
+    usage: Usage
+
+class ErrorResponse(BaseModel):
+    type: str = "error"
+    error: dict
+
+# ============== Helper Functions ==============
+
+def format_messages(messages: List[Message], system: Optional[str] = None) -> str:
+    """Format messages into a prompt string"""
+    formatted_messages = []
+
+    if system:
+        formatted_messages.append({"role": "system", "content": system})
+
+    for msg in messages:
+        content = msg.content
+        if isinstance(content, list):
+            content = " ".join([block.text for block in content if block.type == "text"])
+        formatted_messages.append({"role": msg.role, "content": content})
+
+    # Use chat template if available
+    if tokenizer.chat_template:
+        return tokenizer.apply_chat_template(
+            formatted_messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+
+    # Fallback simple format
+    prompt = ""
+    for msg in formatted_messages:
+        role = msg["role"].capitalize()
+        prompt += f"{role}: {msg['content']}\n"
+    prompt += "Assistant: "
+    return prompt
+
+def generate_id() -> str:
+    """Generate a unique message ID"""
+    return f"msg_{uuid.uuid4().hex[:24]}"
+
+# ============== API Endpoints ==============
+
+@app.get("/")
+async def root():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "model": MODEL_ID,
+        "api_version": "2023-06-01",
+        "compatibility": "anthropic-messages-api"
+    }
+
+@app.get("/v1/models")
+async def list_models():
+    """List available models (Anthropic-compatible)"""
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": "smollm2-135m",
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "huggingface",
+                "display_name": "SmolLM2 135M Instruct"
+            }
+        ]
+    }
+
+@app.post("/v1/messages")
+async def create_message(
+    request: MessageRequest,
+    x_api_key: Optional[str] = Header(None, alias="x-api-key"),
+    anthropic_version: Optional[str] = Header(None, alias="anthropic-version")
+):
+    """
+    Create a message (Anthropic Messages API compatible)
+    """
+    try:
+        # Format the prompt
+        prompt = format_messages(request.messages, request.system)
+
+        # Tokenize
+        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+        input_token_count = inputs.input_ids.shape[1]
+
+        if request.stream:
+            return await stream_response(request, inputs, input_token_count)
+
+        # Generate
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=request.max_tokens,
+                temperature=request.temperature if request.temperature > 0 else 1.0,
+                top_p=request.top_p,
+                top_k=request.top_k,
+                do_sample=request.temperature > 0,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+
+        # Decode only new tokens
+        generated_tokens = outputs[0][input_token_count:]
+        generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+        output_token_count = len(generated_tokens)
+
+        # Build response
+        response = MessageResponse(
+            id=generate_id(),
+            content=[ContentBlock(type="text", text=generated_text.strip())],
+            model=request.model,
+            stop_reason="end_turn",
+            usage=Usage(
+                input_tokens=input_token_count,
+                output_tokens=output_token_count
+            )
+        )
+
+        return response
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+async def stream_response(request: MessageRequest, inputs, input_token_count: int):
+    """Stream response using SSE (Server-Sent Events)"""
+
+    message_id = generate_id()
+
+    async def generate():
+        # Send message_start event
+        start_event = {
+            "type": "message_start",
+            "message": {
+                "id": message_id,
+                "type": "message",
+                "role": "assistant",
+                "content": [],
+                "model": request.model,
+                "stop_reason": None,
+                "stop_sequence": None,
+                "usage": {"input_tokens": input_token_count, "output_tokens": 0}
+            }
+        }
+        yield f"event: message_start\ndata: {json.dumps(start_event)}\n\n"
+
+        # Send content_block_start
+        block_start = {
+            "type": "content_block_start",
+            "index": 0,
+            "content_block": {"type": "text", "text": ""}
+        }
+        yield f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n"
+
+        # Setup streamer
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+        generation_kwargs = {
+            **inputs,
+            "max_new_tokens": request.max_tokens,
+            "temperature": request.temperature if request.temperature > 0 else 1.0,
+            "top_p": request.top_p,
+            "top_k": request.top_k,
+            "do_sample": request.temperature > 0,
+            "pad_token_id": tokenizer.eos_token_id,
+            "eos_token_id": tokenizer.eos_token_id,
+            "streamer": streamer,
+        }
+
+        # Run generation in a thread
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        output_tokens = 0
+        for text in streamer:
+            if text:
+                output_tokens += len(tokenizer.encode(text, add_special_tokens=False))
+                delta_event = {
+                    "type": "content_block_delta",
+                    "index": 0,
+                    "delta": {"type": "text_delta", "text": text}
+                }
+                yield f"event: content_block_delta\ndata: {json.dumps(delta_event)}\n\n"
+
+        thread.join()
+
+        # Send content_block_stop
+        block_stop = {"type": "content_block_stop", "index": 0}
+        yield f"event: content_block_stop\ndata: {json.dumps(block_stop)}\n\n"
+
+        # Send message_delta
+        delta = {
+            "type": "message_delta",
+            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "usage": {"output_tokens": output_tokens}
+        }
+        yield f"event: message_delta\ndata: {json.dumps(delta)}\n\n"
+
+        # Send message_stop
+        yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"
+
+    return StreamingResponse(
+        generate(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no"
+        }
+    )
+
+# Token counting endpoint
+@app.post("/v1/messages/count_tokens")
+async def count_tokens(request: MessageRequest):
+    """Count tokens for a message request"""
+    prompt = format_messages(request.messages, request.system)
+    tokens = tokenizer.encode(prompt)
+    return {"input_tokens": len(tokens)}
+
+# Health check
+@app.get("/health")
+async def health():
+    return {"status": "ok", "model_loaded": model is not None}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)