Spaces:
Sleeping
Sleeping
File size: 5,018 Bytes
55f7a10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import logging
import os
from typing import Optional
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI(
title="DeepSeek R1 Chat API",
description="DeepSeek R1 model hosted on Hugging Face Spaces",
version="1.0.0"
)
# Request/Response models
class ChatRequest(BaseModel):
message: str
max_length: Optional[int] = 512
temperature: Optional[float] = 0.7
top_p: Optional[float] = 0.9
class ChatResponse(BaseModel):
response: str
status: str
# Global variables for model and tokenizer
model = None
tokenizer = None
@app.on_event("startup")
async def load_model():
"""Load the DeepSeek model on startup"""
global model, tokenizer
try:
logger.info("Loading DeepSeek R1 model...")
# Use a smaller DeepSeek model that fits in Spaces
model_name = "deepseek-ai/deepseek-r1-distill-qwen-1.5b"
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(
model_name,
trust_remote_code=True,
padding_side="left"
)
# Add pad token if it doesn't exist
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# Load model with appropriate settings for Spaces
model = AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto" if torch.cuda.is_available() else None,
low_cpu_mem_usage=True
)
logger.info("Model loaded successfully!")
except Exception as e:
logger.error(f"Error loading model: {str(e)}")
raise e
@app.get("/")
async def root():
"""Health check endpoint"""
return {
"message": "DeepSeek R1 Chat API is running!",
"status": "healthy",
"model_loaded": model is not None
}
@app.get("/health")
async def health_check():
"""Detailed health check"""
return {
"status": "healthy",
"model_loaded": model is not None,
"tokenizer_loaded": tokenizer is not None,
"cuda_available": torch.cuda.is_available(),
"device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0
}
@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
"""Chat endpoint for DeepSeek model"""
if model is None or tokenizer is None:
raise HTTPException(status_code=503, detail="Model not loaded yet")
try:
# Prepare the input
prompt = f"User: {request.message}\nAssistant:"
# Tokenize input
inputs = tokenizer(
prompt,
return_tensors="pt",
padding=True,
truncation=True,
max_length=1024
)
# Move to appropriate device
if torch.cuda.is_available():
inputs = {k: v.cuda() for k, v in inputs.items()}
# Generate response
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=request.max_length,
temperature=request.temperature,
top_p=request.top_p,
do_sample=True,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id,
repetition_penalty=1.1
)
# Decode response
full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Extract only the assistant's response
if "Assistant:" in full_response:
response = full_response.split("Assistant:")[-1].strip()
else:
response = full_response[len(prompt):].strip()
return ChatResponse(response=response, status="success")
except Exception as e:
logger.error(f"Error during generation: {str(e)}")
raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
@app.post("/generate")
async def generate(request: ChatRequest):
"""Alternative generation endpoint"""
return await chat(request)
@app.get("/model-info")
async def model_info():
"""Get model information"""
if model is None:
return {"status": "Model not loaded"}
return {
"model_name": "deepseek-ai/deepseek-r1-distill-qwen-1.5b",
"model_type": type(model).__name__,
"tokenizer_type": type(tokenizer).__name__,
"vocab_size": tokenizer.vocab_size if tokenizer else None,
"device": str(next(model.parameters()).device) if model else None
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860) |