jarias commited on
Commit
ba0d2ea
1 Parent(s): 7a9bf02

Upload main.py

Browse files
Files changed (1) hide show
  1. main.py +27 -0
main.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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("zephyr-7b-beta.Q4_K_S.gguf",
7
+ model_type='mistral',
8
+ max_new_tokens = 1096,
9
+ threads = 3,
10
+ )
11
+
12
+
13
+ #Pydantic object
14
+ class validation(BaseModel):
15
+ prompt: str
16
+
17
+ #Fast API
18
+ app = FastAPI()
19
+
20
+ #Zephyr completion
21
+ @app.post("/llm_on_cpu")
22
+ async def stream(item: validation):
23
+ system_prompt = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.'
24
+ E_INST = "</s>"
25
+ user, assistant = "<|user|>", "<|assistant|>"
26
+ prompt = f"{system_prompt}{E_INST}\n{user}\n{item.prompt.strip()}{E_INST}\n{assistant}\n"
27
+ return llm(prompt)