import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. Model ve Cihaz Ayarları
MODEL_NAME = "AhiskaAI/AhiskaAI-134m-Instruct-v0.2"
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

print(f"[AhıskaAI] Model yükleniyor: {MODEL_NAME}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, 
    torch_dtype=dtype,
    device_map="auto" if torch.cuda.is_available() else None
).to(device)
model.eval()

# Stop Token Yönetimi
try:
    stop_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
except:
    stop_token_id = tokenizer.eos_token_id

# 2. Çıkarım Fonksiyonu (Stateless / Tek Turlu)
def predict(message, history):
    SYSTEM_PROMPT = "Sen kibar, sorulan soruları tam cümlelerle yanıtlayan Türkçe bir asistansın."
    
    prompt = (
        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
        f"<|im_start|>user\n{message}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=128,          
            do_sample=True,            
            temperature=0.3,           
            top_p=0.9,                  
            repetition_penalty=1.25,    
            no_repeat_ngram_size=3,     
            eos_token_id=stop_token_id, 
            pad_token_id=tokenizer.pad_token_id
        )
    
    generated_ids = outputs[0][len(inputs.input_ids[0]):]
    response = tokenizer.decode(generated_ids, skip_special_tokens=False)
    
    # Temizleme
    clean_response = response.replace("<|im_end|>", "").replace("<|im_start|>", "").strip()
    return clean_response

# 3. Profesyonel Arayüz Tasarımı
with gr.Blocks(title="AhıskaAI Araştırma Laboratuvarı") as demo:
    gr.Markdown(
        """
        # AhıskaAI-134m-Instruct-v0.2
        **AhıskaAI Laboratuvarı tarafından sıfırdan eğitilmiş; yüksek kaliteli, Türkçe diline ve kültürel bağlama odaklı kompakt yapay zeka asistanı.**
        
        * **Mimari:** Llama tabanlı, 134 Milyon parametre.
        * **Eğitim:** 1.2GB filtrelenmiş Türkçe veri seti (SFT).
        * **Mod:** Stateless (Hafızasız/Tek Turlu) - Her soru bağımsız işlenir.
        """
    )
    
    gr.ChatInterface(fn=predict)
    
    gr.Markdown(
        """
        ---
        **Not:** Bu, deneysel bir modeldir. Yanıtlar sınırlı parametre kapasitesiyle üretilir.
        """
    )

if __name__ == "__main__":
    demo.launch()