from flask import Flask, request, jsonify
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os

app = Flask(__name__)

# Free Tier එක සඳහා 2B මාදිලිය වඩාත් සුදුසුයි
MODEL_ID = "google/gemma-2-2b-it" 

print("Loading Olya.ai... Please wait ✨")

# Hugging Face Token එක ලබා ගැනීම
token = os.environ.get("HF_TOKEN")

# Tokenizer සහ Model එක Load කිරීම
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=token)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    torch_dtype=torch.float16, # RAM එක ඉතිරි කර ගැනීමට float16 භාවිතා කරමු
    low_cpu_mem_usage=True,
    token=token
)

system_instructions = (
    "You are Olya.ai, a highly intelligent and elegant AI assistant developed by Hansaka P Fernando. "
    "Your primary language is Sinhala. Respond in clear, polite, and warm Sinhala. "
    "Use 'ඔයා' (Oya) when addressing the user."
)

@app.route('/')
def home():
    try:
        with open("index.html", "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        return "Error: index.html not found.", 404

@app.route('/api/chat', methods=['POST'])
def chat():
    data = request.json
    user_message = data.get("message", "")
    
    if not user_message:
        return jsonify({"reply": "කරුණාකර පණිවිඩයක් ඇතුළත් කරන්න."})

    # Gemma සඳහා නිවැරදි Chat Format එක
    messages = [
        {"role": "user", "content": f"{system_instructions}\n\n{user_message}"},
    ]
    
    # මෙහිදී return_dict=True සහ **inputs භාවිතා කිරීමෙන් දෝෂය මගහැරේ
    inputs = tokenizer.apply_chat_template(
        messages, 
        add_generation_prompt=True, 
        return_tensors="pt",
        return_dict=True
    ).to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs, # මෙහිදී දත්ත නිවැරදිව unpack කර ලබා දෙයි
            max_new_tokens=256,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.2
        )
    
    # පිළිතුර පමණක් වෙන් කර ගැනීම
    input_length = inputs.input_ids.shape[1]
    response_text = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
    
    return jsonify({"reply": response_text.strip()})

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