hArshi07 commited on
Commit
a316563
1 Parent(s): aa029c2
Files changed (1) hide show
  1. main.py +21 -0
main.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ctransformers import AutoModelForCausalLM
2
+ from fastapi import FastAPI
3
+ from pydantic import BaseModel
4
+
5
+ llm = AutoModelForCausalLM.from_pretrained("TheBloke/CodeLlama-7B-Instruct-GGUF",
6
+ model_file="codellama-7b-instruct.q4_K_M.gguf",
7
+ model_type="llama",
8
+ gpu_layers=0)
9
+ #Pydantic object
10
+ class validation(BaseModel):
11
+ prompt: str
12
+ #Fast API
13
+ app = FastAPI()
14
+
15
+ @app.post("/llm_on_cpu")
16
+ async def stream(item: validation):
17
+ system_prompt = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.'
18
+ E_INST = "</s>"
19
+ user, assistant = "<|user|>", "<|assistant|>"
20
+ prompt = f"{system_prompt}{E_INST}\n{user}\n{item.prompt}{E_INST}\n{assistant}\n"
21
+ return llm(prompt)