import os
import time
from flask import Flask, request, render_template
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
import torch

app = Flask(__name__)

BASE_MODEL = "EleutherAI/gpt-neox-20b"
LORA_ADAPTER = "./model"

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# 4-bit quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

print("Loading 4-bit base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    quantization_config=bnb_config,
    device_map="auto"
)

print("Loading LoRA adapter...")
model = PeftModel.from_pretrained(
    base_model,
    LORA_ADAPTER,
    device_map="auto"
)

model.eval()
print("Model loaded successfully in 4-bit!")

# -----------------------------
# Flask Routes
# -----------------------------
@app.route("/", methods=["GET", "POST"])
def index():
    output_text = ""
    if request.method == "POST":
        user_input = request.form.get("user_input")
        if user_input:
            inputs = tokenizer(user_input, return_tensors="pt").to(device)
            with torch.no_grad():
                generated_ids = model.generate(
                    **inputs,
                    max_new_tokens=200,
                    do_sample=True,
                    temperature=0.7
                )
            output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return render_template("index.html", output_text=output_text)

# Give HF time to detect the server
time.sleep(2)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))