import torch
import gradio as gr
import spaces
import re # Metin temizleme için eklendi
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# --- Model Ayarları ---
base_model_name = "unsloth/gpt-oss-20b"
adapter_model_name = "userdotcs/gpt-oss-20b-turkish-correction-adapter"

print("Model yükleniyor...")
tokenizer = AutoTokenizer.from_pretrained(base_model_name)

base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

model = PeftModel.from_pretrained(base_model, adapter_model_name)
model.eval()

@spaces.GPU(duration=120)
def fix_text(input_text):
    if not input_text or input_text.strip() == "":
        return ""

    messages = [
        {
            "role": "system",
            "content": "You are an intelligent assistant that corrects Turkish spelling and grammar mistakes."
        },
        {
            "role": "user",
            "content": f"Fix typos in the text:\n{input_text}"
        }
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs,
            max_new_tokens=16_384,
            pad_token_id=tokenizer.eos_token_id
        )

    input_length = inputs.shape[1]
    full_response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)

    # --- Ayıklama Mantığı ---
    # Modelin çıktısında "assistantfinal" anahtar kelimesini arıyoruz
    separator = "assistantfinal"
    
    if separator in full_response:
        # Metni separator'a göre böl ve en son (sağdaki) parçayı al
        clean_response = full_response.split(separator)[-1]
    else:
        # Eğer separator yoksa ama "analysis" varsa, analizi temizlemeye çalışalım
        # Genellikle asıl cevap en sonda olur. 
        # Alternatif olarak son satırı almayı deneyebiliriz:
        lines = full_response.strip().split('\n')
        clean_response = lines[-1] if lines else full_response

    return clean_response.strip()

# Arayüz
demo = gr.Interface(
    fn=fix_text,
    inputs=gr.Textbox(label="Input", lines=3),
    outputs=gr.Textbox(label="Output", lines=3),
    title="gpt-oss-20b Turkish correction"
)

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