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

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +10 -11
main.py CHANGED
@@ -2,23 +2,22 @@ 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)
 
2
  from fastapi import FastAPI
3
  from pydantic import BaseModel
4
 
 
5
  llm = AutoModelForCausalLM.from_pretrained("zephyr-7b-beta.Q4_K_S.gguf",
6
+ model_type='mistral',
7
+ max_new_tokens=1096,
8
+ threads=3)
 
9
 
10
+ # Pydantic object
11
+ class Validation(BaseModel):
12
  prompt: str
13
+
14
+ # Fast API
15
  app = FastAPI()
16
 
17
  @app.post("/llm_on_cpu")
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 = "", ""
22
  prompt = f"{system_prompt}{E_INST}\n{user}\n{item.prompt}{E_INST}\n{assistant}\n"
23
+ return await llm(prompt) # Making the call asynchronous