app.py

import os
import time
from flask import Flask, jsonify, request
from flask_cors import CORS 
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig  

# Set the HF_HOME environment variable to a writable directory
os.environ["HF_HOME"] = "/workspace/huggingface_cache"  # Change this to a writable path in your space

app = Flask(__name__)

# Enable CORS for specific origins
CORS(app, resources={r"/api/predict/*": {"origins": ["http://localhost:3000", "https://main.dbn2ikif9ou3g.amplifyapp.com"]}})

# Global variables for model and tokenizer
model = None
tokenizer = None
last_loaded_time = 0
COOLDOWN_PERIOD = 300  # Set your cooldown period to 5 minutes (300 seconds)

def get_model_and_tokenizer(model_id):
    global model, tokenizer
    try:
        print(f"Loading tokenizer for model_id: {model_id}")
        # Load the tokenizer
        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
        tokenizer.pad_token = tokenizer.eos_token
        
        print(f"Loading model for model_id: {model_id}")
        # Load the model
        model = AutoModelForCausalLM.from_pretrained(model_id)  # , device_map="auto")
        model.config.use_cache = False
        print("Model loaded successfully!")
    except Exception as e:
        print(f"Error loading model: {e}")

def is_model_loaded_and_fresh():
    global last_loaded_time
    current_time = time.time()
    return model is not None and (current_time - last_loaded_time) < COOLDOWN_PERIOD

def generate_response(user_input, model_id):
    prompt = formatted_prompt(user_input)
    
    global model, tokenizer

    # Check if model is loaded and fresh
    if not is_model_loaded_and_fresh():
        get_model_and_tokenizer(model_id)  # Load model and tokenizer
        global last_loaded_time
        last_loaded_time = time.time()  # Update the last load time

    # Prepare the input tensors
    inputs = tokenizer(prompt, return_tensors="pt")  # Move inputs to GPU if available
    
    generation_config = GenerationConfig(
        max_new_tokens=100,
        min_length=5,
        do_sample=False,
        num_beams=1,
        pad_token_id=tokenizer.eos_token_id,
        truncation=True
    )

    try:
        # Generate response
        outputs = model.generate(**inputs, generation_config=generation_config)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response
    except Exception as e:
        print(f"Error generating response: {e}")
        return "Error generating response."
    
def formatted_prompt(question) -> str:
    return f"<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant:"

@app.route("/", methods=["GET"])
def handle_get_request():
    message = request.args.get("message", "No message provided.")
    return jsonify({"message": message, "status": "GET request successful!"})

@app.route("/send_message", methods=["POST"])
def handle_post_request():
    data = request.get_json()
    if data is None:
        return jsonify({"error": "No JSON data provided"}), 400

    message = data.get("inputs", "No message provided.") 
    model_id = data.get("model_id", "YALCINKAYA/FinetunedByYalcin")  # Default model if not provided

    try:
        # Generate a response from the model
        model_response = generate_response(message, model_id)
        return jsonify({
            "received_message": model_response, 
            "model_id": model_id, 
            "status": "POST request successful!"
        })
    except Exception as e:
        print(f"Error handling POST request: {e}")
        return jsonify({"error": "An error occurred while processing your request."}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)