Update main.py

main.py

@@ -2,6 +2,8 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 import requests
 from llama_cpp import Llama 
+import threading
+import gc
 
 llms = {
   "TinyLLama 1b 4_K_M 2048": {
@@ -39,6 +41,12 @@ llms = {
     
 #Fast API
 app = FastAPI()
+llm = None
+
+def clean_memory():
+    llm = None
+    gc.collect()
+    
 
 @app.post("/llm_on_cpu")
 async def stream(item: dict):
@@ -54,4 +62,7 @@ async def stream(item: dict):
     llm = Llama(model_path="./code/"+model['file'], n_ctx=nctx, verbose=True,  n_threads=8)
 
     prompt = f"{prefix}{user.replace('{prompt}', item['prompt'])}{suffix}"
-    return llm(prompt, max_tokens=max_tokens)
+    result = llm(prompt, max_tokens=max_tokens)
+    thread = threading.Thread(target=clean_memory)
+    thread.start()
+    return result