Update main.py

main.py

@@ -4,34 +4,64 @@ import requests
 from ctransformers import AutoModelForCausalLM
 
 llms = {
-"tinyllama":{"name": "TheBloke/TinyLlama-1.1B-1T-OpenOrca-GGUF", "file":"tinyllama-1.1b-1t-openorca.Q4_K_M.gguf", "suffix":"<|im_end|><|im_start|>assistant", "prefix":"<|im_start|>system You are a helpful assistant  <|im_end|><|im_start|>user"},
-"tinyllama2":{"name": "TheBloke/TinyLlama-1.1B-1T-OpenOrca-GGUF", "file":"tinyllama-1.1b-1t-openorca.Q3_K_M.gguf", "suffix":"<|im_end|><|im_start|>assistant", "prefix":"<|im_start|>system You are a helpful assistant  <|im_end|><|im_start|>user"},
-"tinyllama3":{"name": "TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF", "file":"tinyllama-1.1b-chat-v0.3.Q4_K_M.gguf", "suffix":"<|im_end|><|im_start|>assistant", "prefix":"<|im_start|>system You are a helpful assistant  <|im_end|><|im_start|>user"},
-"tinyllama4":{"name": "TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF", "file":"tinyllama-1.1b-chat-v0.3.Q2_K.gguf", "suffix":"<|im_end|><|im_start|>assistant", "prefix":"<|im_start|>system You are a helpful assistant  <|im_end|><|im_start|>user"},
-"solar":{"name": "TheBloke/SOLAR-10.7B-Instruct-v1.0-GGUF", "file":"solar-10.7b-instruct-v1.0.Q4_K_M.gguf", "suffix":"\n### Assistant:\n", "prefix":"### User:\n"},
-#"mixtral-moe":{"name": "TheBloke/Mixtral_7Bx2_MoE-GGUF", "file":"mixtral_7bx2_moe.Q4_K_M.gguf", "suffix":"\n### Assistant:\n", "prefix":"### User:\n"},
-#"phi2":{"name": "TheBloke/phi-2-GGUF", "file":"phi-2.Q4_K_M.gguf", "suffix":"", "prefix":""}
+  "TinyLLama 1b 4_K_M 2048": {
+    "nctx": 2048, 
+    "file": "tinyllama-1.1b-chat-v0.3.Q4_K_M.gguf", 
+    "prefix": "### Human:", 
+    "suffix": "### Assistant:"
+  },
+  "TinyLLama 1b OpenOrca 4_K_M 2048": {
+    "nctx": 2048, 
+    "file": "tinyllama-1.1b-1t-openorca.Q4_K_M.gguf", 
+    "prefix": "<|im_start|>system You are a helpfull assistant<|im_end|><|im_start|>user", 
+    "suffix": "<|im_end|><|im_start|>assistant"
+  },
+  "OpenLLama 3b 4_K_M 196k": {
+    "nctx": 80000, 
+    "file": "open-llama-3b-v2-wizard-evol-instuct-v2-196k.Q4_K_M.gguf", 
+    "prefix": "### HUMAN:", 
+    "suffix": "### RESPONSE:"
+  },
+  "Phi-2 2.7b 4_K_M 2048": {
+    "nctx": 2048, 
+    "file": "phi-2.Q4_K_M.gguf", 
+    "prefix": "Instruct:", 
+    "suffix": "Output:"
+  },
+  "Mixtral MOE 7bx2 4_K_M 32K": {
+    "nctx": 32000, 
+    "file": "mixtral_7bx2_moe.Q4_K_M.gguf", 
+    "prefix": "", 
+    "suffix": ""
+  },
+  "Stable Zephyr 3b 4_K_M 4096": {
+    "nctx": 4096, 
+    "file": "stablelm-zephyr-3b.Q4_K_M.gguf", 
+    "prefix": "<|user|>", 
+    "suffix": "<|endoftext|><|assistant|>"
+  }
 }
 
-for k in llms.keys():
-    AutoModelForCausalLM.from_pretrained(llms[k]['name'], model_file=llms[k]['file'])
-
 #Pydantic object
 class validation(BaseModel):
     prompt: str
     llm: str
+    
 #Fast API
 app = FastAPI()
 
 @app.post("/llm_on_cpu")
 async def stream(item: validation):
 
-    prefix=llms[item.llm]['prefix']
-    suffix=llms[item.llm]['suffix']
+    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
     user="""
     {prompt}"""
     
-    llm = AutoModelForCausalLM.from_pretrained(llms[item.llm]['name'], model_file=llms[item.llm]['file'])
+    model = Llama(model_path="./"+model['file'], n_ctx=model['nctx'], verbose=False,  n_threads=8)
 
     prompt = f"{prefix}{user.replace('{prompt}', item.prompt)}{suffix}"
-    return llm(prompt)
+    return llm(prompt, max_tokens=max_tokens)