Update app.py

app.py

@@ -1,28 +1,92 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-import uvicorn
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import logging
+from typing import List
+import os
+import uuid
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-vl-1.3b-chat")
-model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-vl-1.3b-chat")
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Initialize FastAPI app
-app = FastAPI()
+app = FastAPI(title="DeepSeek CPU Hosting API")
+
+# Model configuration
+MODEL_NAME = "deepseek-ai/deepseek-llm-7b"  # Example model, replace with actual DeepSeek model
+DEVICE = "cpu"  # Force CPU usage
+
+# Load model and tokenizer
+try:
+    logger.info("Loading model and tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    model.to(DEVICE)
+    logger.info("Model loaded successfully!")
+except Exception as e:
+    logger.error(f"Failed to load model: {str(e)}")
+    raise
+
+# API key storage (in production, use a proper database)
+API_KEYS = {}
 
-# Define input schema
-class RequestBody(BaseModel):
+# Request models
+class GenerationRequest(BaseModel):
     prompt: str
+    max_length: int = 100
+    temperature: float = 0.7
+    top_p: float = 0.9
+
+class APIKeyRequest(BaseModel):
+    name: str
+
+# Generation endpoint
+@app.post("/generate/{api_key}")
+async def generate_text(api_key: str, request: GenerationRequest):
+    if api_key not in API_KEYS:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+    
+    try:
+        inputs = tokenizer(request.prompt, return_tensors="pt").to(DEVICE)
+        outputs = model.generate(
+            **inputs,
+            max_length=request.max_length,
+            temperature=request.temperature,
+            top_p=request.top_p,
+            do_sample=True
+        )
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        logger.info(f"Generated text for API key: {api_key}")
+        return {"generated_text": generated_text}
+    except Exception as e:
+        logger.error(f"Generation error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# API key management endpoints
+@app.post("/create_api_key")
+async def create_api_key(request: APIKeyRequest):
+    new_key = str(uuid.uuid4())
+    API_KEYS[new_key] = {
+        "name": request.name,
+        "usage_count": 0
+    }
+    logger.info(f"Created new API key for {request.name}")
+    return {"api_key": new_key, "name": request.name}
+
+@app.get("/list_api_keys")
+async def list_api_keys():
+    return {"api_keys": API_KEYS}
 
-# Define the model inference function
-@app.post("/predict")
-async def predict(request: RequestBody):
-    inputs = tokenizer(request.prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"response": result}
+@app.delete("/revoke_api_key/{api_key}")
+async def revoke_api_key(api_key: str):
+    if api_key in API_KEYS:
+        del API_KEYS[api_key]
+        logger.info(f"Revoked API key: {api_key}")
+        return {"status": "success"}
+    raise HTTPException(status_code=404, detail="API key not found")
 
-# For testing locally (not needed for Hugging Face)
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)