File size: 3,832 Bytes
c1f5657 | 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 | #!/usr/bin/env python3
# ============================================================
# ZabaanAI-v2: FastAPI Server for Production Deployment
# Run: uvicorn scripts.06_deploy_api:app --host 0.0.0.0 --port 8000
# ============================================================
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, List
import torch
import os, warnings
warnings.filterwarnings('ignore')
app = FastAPI(title='ZabaanAI v2 API', version='2.0.0',
description='Pakistan Multilingual AI - Supports Urdu, Punjabi, Sindhi, Pashto, Balochi, Saraiki, English, Roman Urdu')
USE_GPU = os.getenv('USE_GPU', '0') == '1'
MODEL_PATH = os.getenv('MODEL_PATH', 'shaikhsalman/zabaanai-v2-sft')
DEVICE = 'cuda' if USE_GPU else 'cpu'
# Load on startup
print(f'Loading model from {MODEL_PATH} on {DEVICE}...')
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
base = AutoModelForCausalLM.from_pretrained(
'Qwen/Qwen2.5-7B-Instruct',
device_map=DEVICE,
load_in_4bit=(not USE_GPU),
torch_dtype=torch.bfloat16 if USE_GPU else torch.float16,
trust_remote_code=True
)
model = PeftModel.from_pretrained(base, MODEL_PATH)
model.eval()
print('Model ready!')
SYSTEM_MSG = (
'You are ZabaanAI, a helpful multilingual AI assistant specializing in '
'Pakistan languages. Respond in the same language as the user.'
)
class Message(BaseModel):
role: str
content: str
class ChatRequest(BaseModel):
messages: List[Message]
max_tokens: int = 512
temperature: float = 0.7
top_p: float = 0.9
class GenerateRequest(BaseModel):
prompt: str
max_tokens: int = 512
temperature: float = 0.7
system: Optional[str] = None
def build_prompt(messages: List[Message], system: str = None) -> str:
sys_msg = system or SYSTEM_MSG
prompt = f'<|im_start|>system\n{sys_msg}<|im_end|>\n'
for m in messages:
role = 'user' if m.role in ('user', 'human') else 'assistant'
prompt += f'<|im_start|>{role}\n{m.content}<|im_end|>\n'
prompt += '<|im_start|>assistant\n'
return prompt
@app.post('/v1/chat/completions')
async def chat_completions(req: ChatRequest):
prompt = build_prompt(req.messages)
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=req.max_tokens,
temperature=req.temperature,
do_sample=True,
top_p=req.top_p,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
return {'model': 'zabaanai-v2', 'choices': [{'message': {'role': 'assistant', 'content': response.strip()}}]}
@app.post('/v1/generate')
async def generate(req: GenerateRequest):
sys_msg = req.system or SYSTEM_MSG
prompt = f'<|im_start|>system\n{sys_msg}<|im_end|>\n<|im_start|>user\n{req.prompt}<|im_end|>\n<|im_start|>assistant\n'
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs, max_new_tokens=req.max_tokens,
temperature=req.temperature, do_sample=True,
pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
)
return {'text': tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True).strip()}
@app.get('/health')
async def health():
return {'status': 'ok', 'model': 'zabaanai-v2'}
if __name__ == '__main__':
import uvicorn
uvicorn.run(app, host='0.0.0.0', port=8000)
|