SrejonAhamed commited on
Commit
f71eea5
1 Parent(s): de3fa38

Upload main.py

Browse files
Files changed (1) hide show
  1. main.py +24 -0
main.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ctransformers import AutoModelForCausalLM
2
+ from fastapi import FastAPI
3
+ from pydantic import BaseModel
4
+
5
+
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
+ #Pydantic object
13
+ class validation(BaseModel):
14
+ prompt: str
15
+ #Fast API
16
+ app = FastAPI()
17
+
18
+ @app.post("/llm_on_cpu")
19
+ async def stream(item: validation):
20
+ system_prompt = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.'
21
+ E_INST = "</s>"
22
+ user, assistant = "<|user|>", "<|assistant|>"
23
+ prompt = f"{system_prompt}{E_INST}\n{user}\n{item.prompt}{E_INST}\n{assistant}\n"
24
+ return llm(prompt)