File size: 1,441 Bytes
624ccff 3ca9fbc 2d9bd99 53bcd1d 2d9bd99 624ccff 3ca9fbc 624ccff |
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 |
from ctransformers import AutoModelForCausalLM
from fastapi import FastAPI, Form
from pydantic import BaseModel
#Model loading
llm = AutoModelForCausalLM.from_pretrained("TheBloke/Toppy-M-7B-GGUF",
model_file="toppy-m-7b.Q5_K_M.gguf",
model_type="mistral",
context_length=4096,
temperature=1,
gpu_layers=50)
# llm = AutoModelForCausalLM.from_pretrained("TheDrummer/Moistral-11B-v3-GGUF",
# model_file="Moistral-11B-v3-Q6_K.gguf",
# model_type="mistral",
# context_length=4096,
# temperature=0.6,
# gpu_layers=50)
#Pydantic object
class validation(BaseModel):
prompt: str
#Fast API
app = FastAPI()
#Zephyr completion
@app.post("/llm_on_gpu")
async def stream(item: validation):
system_prompt = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.'
E_INST = "</s>"
user, assistant = "<|user|>", "<|assistant|>"
prompt = f"{system_prompt}{E_INST}\n{user}\n{item.prompt.strip()}{E_INST}\n{assistant}\n"
return llm(prompt) |