Xeltron-cloud/VexaAI_DeepSeek-Coder-6.7B-Instruct_API

1from fastapi import FastAPI, HTTPException2from pydantic import BaseModel3from transformers import AutoTokenizer, AutoModelForCausalLM4from huggingface_hub import login5import os6import torch7import uvicorn8 9login(os.getenv("HF_TOKEN"))10 11app = FastAPI(12    title="VexaAI Model-Platform: DeepSeek-Coder-6.7B-Instruct",13    description="Self-hosted AI-Model DeepSeek-Coder-6.7B-Instruct, powered by VexaAI.",14    version="0.9"15)16 17model_name = "TheBloke/deepseek-coder-6.7B-instruct-GGUF"18 19tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)20model = AutoModelForCausalLM.from_pretrained(21    model_name,22    device_map="auto",23    trust_remote_code=True,24    torch_dtype=torch.float3225)26model.eval()27 28class GenerateRequest(BaseModel):29    prompt: str30    max_new_tokens: int = 51231    temperature: float = 0.732 33@app.post("/generate")34async def generate_text(request: GenerateRequest):35    try:36        inputs = tokenizer(request.prompt, return_tensors="pt").to(model.device)37        38        with torch.no_grad():39            outputs = model.generate(40                **inputs,41                max_new_tokens=request.max_new_tokens,42                temperature=request.temperature,43                do_sample=True,44                repetition_penalty=1.1,45                eos_token_id=tokenizer.eos_token_id,46                pad_token_id=tokenizer.eos_token_id47            )48        49        full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)50        generated_text = full_text[len(tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True)):].strip()51        52        return {"generated_text": generated_text}53    except Exception as e:54        raise HTTPException(status_code=500, detail=f"VexaAI-Lab: HTTP/S error: {str(e)}")55 56@app.get("/")57async def root():58    return {"message": "To start generating text, use /generate."}59 60if __name__ == "__main__":61    uvicorn.run(app, host="0.0.0.0", port=7860)