Update main.py

main.py

@@ -46,6 +46,9 @@ llms = {
 class validation(BaseModel):
     prompt: str
     llm: str
+    max_tokens: int
+    nctx: int
+    
     
 #Fast API
 app = FastAPI()
@@ -56,8 +59,8 @@ async def stream(item: validation):
     model = llms[item.llm]
     prefix=model['prefix']
     suffix=model['suffix']
-    nctx = model['nctx'] if 'nctx' in item.keys() else 1024
-    max_tokens = model['max_tokens'] if 'max_tokens' in item.keys() else 512
+    nctx = model['nctx'] if item.nctx is not None else 1024
+    max_tokens = model['max_tokens'] if item.max_tokens is not None else 512
     user="""
     {prompt}"""