restructureapp

Dockerfile

@@ -11,4 +11,4 @@ COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
 COPY --chown=user . /app
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "app:main", "--host", "0.0.0.0", "--port", "7860"]

app.py → app/main.py

@@ -32,10 +32,15 @@ tensor = torch.randn(3, 3).to(device)
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     print("Starting the app")   
-    yield
-    print("Stopping the app")
+    try:
+        yield
+    finally:
+        print("Stopping the app")
+        # Ensure proper cleanup
+        del tensor
+        torch.cuda.empty_cache()
 
-app = FastAPI()
+app = FastAPI(lifespan=lifespan)
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, token=huggingface_token)
 model = AutoModelForCausalLM.from_pretrained(model_name, token=huggingface_token)
@@ -88,7 +93,7 @@ async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
     await websocket.send_text("Hello, this is the Heisterbot chat server. Please send a message /chat")
 
-@app.websocket("/chat", name="chat", summary="Chat with the Heisterbot", description="Chat with the Heisterbot")
+@app.websocket("/chat",)
 async def websocket_endpoint(websocket: WebSocket):
     print("Client connected at: ", websocket.client)
     await websocket.accept()