Xeltron-cloud commited on
Commit
8b5d331
·
verified ·
1 Parent(s): fa2fa15

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +61 -0
app.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, HTTPException
2
+ from pydantic import BaseModel
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+ from huggingface_hub import login
5
+ import os
6
+ import torch
7
+ import uvicorn
8
+
9
+ login(os.getenv("HF_TOKEN"))
10
+
11
+ app = 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
+
17
+ model_name = "TheBloke/deepseek-coder-6.7B-instruct-GGUF"
18
+
19
+ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
20
+ model = AutoModelForCausalLM.from_pretrained(
21
+ model_name,
22
+ device_map="auto",
23
+ trust_remote_code=True,
24
+ torch_dtype=torch.float32
25
+ )
26
+ model.eval()
27
+
28
+ class GenerateRequest(BaseModel):
29
+ prompt: str
30
+ max_new_tokens: int = 512
31
+ temperature: float = 0.7
32
+
33
+ @app.post("/generate")
34
+ async 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_id
47
+ )
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("/")
57
+ async def root():
58
+ return {"message": "To start generating text, use /generate."}
59
+
60
+ if __name__ == "__main__":
61
+ uvicorn.run(app, host="0.0.0.0", port=7860)