wizardcoder-ggml / main.py
vasilee's picture
Update main.py
3481f63
raw
history blame contribute delete
971 Bytes
import fastapi
import uvicorn
from fastapi.middleware.cors import CORSMiddleware
from ctransformers import AutoModelForCausalLM
from pydantic import BaseModel
llm = AutoModelForCausalLM.from_pretrained("TheBloke/WizardLM-7B-V1.0-Uncensored-GGML",
model_file="wizardlm-7b-v1.0-uncensored.ggmlv3.q2_K.bin",
model_type="llama")
app = fastapi.FastAPI(title="WizardLM-7B-V1.0-Uncensored")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/")
async def index():
return {"message": "Hello World"}
class InferenceRequest(BaseModel):
input: str
@app.post("/inference")
async def inference(request: InferenceRequest):
generated_text = llm(request.input)
return {"generated_text": generated_text}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)