Update app.py

app.py

@@ -1,16 +1,18 @@
-# app.py
 import os
 import torch
 from flask import Flask, request, jsonify
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
 
 app = Flask(__name__)
 
-# Load the model and tokenizer
-@app.before_first_request
-def load_model():
+# Global variables to store model, tokenizer, and dataset
+model = None
+tokenizer = None
+dataset = None
+
+# Function to load the model and dataset
+def load_model_and_data():
     global model, tokenizer, dataset
     
     # Load the base model and tokenizer
@@ -28,6 +30,13 @@ def load_model():
     dataset = load_dataset("lahiruchamika27/tia")
     print("Model, tokenizer, and dataset loaded successfully!")
 
+# Initialize the model on startup
+@app.before_request
+def before_request():
+    global model, tokenizer, dataset
+    if model is None or tokenizer is None or dataset is None:
+        load_model_and_data()
+
 # Define chat endpoint
 @app.route('/api/chat', methods=['POST'])
 def chat():