Update main.py

main.py

@@ -41,30 +41,23 @@ llms = {
     "suffix": "<|endoftext|><|assistant|>"
   }
 }
-
-#Pydantic object
-class validation(BaseModel):
-    prompt: str
-    llm: str
-    max_tokens: int = 512
-    nctx: int = 2048
     
     
 #Fast API
 app = FastAPI()
 
 @app.post("/llm_on_cpu")
-async def stream(item: validation):
+async def stream(item: dict):
 
-    model = llms[item.llm]
+    model = llms[item['llm']]
     prefix=model['prefix']
     suffix=model['suffix']
-    nctx =  item.nctx if item.nctx is not None else model['nctx']
-    max_tokens = item.max_tokens if item.max_tokens is not None else 512
+    nctx =  item['nctx'] if 'nctx' is not in item.keys() else model['nctx']
+    max_tokens = item['max_tokens'] if 'max_tokens' is not in item.keys() else 512
     user="""
     {prompt}"""
     
     llm = Llama(model_path="./code/"+model['file'], n_ctx=model['nctx'], verbose=False,  n_threads=8)
 
-    prompt = f"{prefix}{user.replace('{prompt}', item.prompt)}{suffix}"
+    prompt = f"{prefix}{user.replace('{prompt}', item['prompt'])}{suffix}"
     return llm(prompt, max_tokens=max_tokens)