zehaoj commited on
Commit
624ccff
·
verified ·
1 Parent(s): 88fbc17

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +23 -0
main.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ctransformers import AutoModelForCausalLM
2
+ from fastapi import FastAPI, Form
3
+ from pydantic import BaseModel
4
+
5
+ #Model loading
6
+ llm = AutoModelForCausalLM.from_pretrained("TheBloke/Toppy-M-7B-GGUF", model_file="toppy-m-7b.Q5_K_S.gguf", model_type="mistral", gpu_layers=50)
7
+
8
+
9
+ #Pydantic object
10
+ class validation(BaseModel):
11
+ prompt: str
12
+
13
+ #Fast API
14
+ app = FastAPI()
15
+
16
+ #Zephyr completion
17
+ @app.post("/llm_on_gpu")
18
+ async def stream(item: validation):
19
+ system_prompt = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.'
20
+ E_INST = "</s>"
21
+ user, assistant = "<|user|>", "<|assistant|>"
22
+ prompt = f"{system_prompt}{E_INST}\n{user}\n{item.prompt.strip()}{E_INST}\n{assistant}\n"
23
+ return llm(prompt)