import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
)
from datasets import Dataset
import numpy as np
from sklearn.metrics import accuracy_score, f1_score

print(torch.cuda.is_available(), torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU only")

import pandas as pd
from datasets import Dataset, DatasetDict

train_df = pd.read_csv(
    "https://raw.githubusercontent.com/e9t/nsmc/master/ratings_train.txt",
    sep="\t"
).dropna()
test_df = pd.read_csv(
    "https://raw.githubusercontent.com/e9t/nsmc/master/ratings_test.txt",
    sep="\t"
).dropna()

raw_dataset = DatasetDict({
    "train": Dataset.from_pandas(train_df, preserve_index=False),
    "test": Dataset.from_pandas(test_df, preserve_index=False),
})

print(raw_dataset)
print("============")
for i in range(8):
    print(raw_dataset["train"][i]["document"], "==>", raw_dataset["train"][i]["label"])

model_name = "klue/bert-base"  # 한국어 사전학습 BERT
tokenizer = AutoTokenizer.from_pretrained(model_name)

def tokenize_fn(batch):
    return tokenizer(
        batch["document"],
        truncation=True,
        padding=False,
        max_length=64,
    )

tokenized_dataset = raw_dataset.map(tokenize_fn, batched=True)
print(tokenized_dataset["train"][0])

sample = "이 영화 진짜 최고였어요, 다시 보고 싶어요"
encoded = tokenizer(sample)
print("토큰 ID:", encoded["input_ids"])
print("토큰 복원:", tokenizer.convert_ids_to_tokens(encoded["input_ids"]))

from transformers import AutoModel

embed_model = AutoModel.from_pretrained(model_name) 
# 토큰라이저로 만든 일차원리스트중 인덱스값을(encoded["input_ids"]) 토치의 텐서(단[]로처리하여 2차원이됨)로 변경
input_ids = torch.tensor([encoded["input_ids"]])

with torch.no_grad():
    outputs = embed_model(input_ids)

print("임베딩 벡터 shape:", outputs.last_hidden_state.shape)
# (배치=1, 시퀀스 길이, hidden_size=768) → 토큰 하나하나가 768차원 벡터로 표현됨

# 예: 시퀀스에서 3번째 토큰(index=2)의 임베딩 벡터 확인
token_idx = 2

# outputs.last_hidden_state shape: (batch=1, seq_len, hidden_size=768)
vector = outputs.last_hidden_state[0, token_idx]  # shape: (768,)

print("전체 토큰:",input_ids)
print("해당 토큰:", input_ids[0, token_idx], tokenizer.convert_ids_to_tokens([encoded["input_ids"][token_idx]]))
print("벡터 shape:", vector.shape)
print("벡터 값 (앞 10개만):", vector[:10])
print("전체 벡터:\n", vector)

model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=2,
)

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=2,
)

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    return {
        "accuracy": accuracy_score(labels, preds),
        "f1": f1_score(labels, preds),
    }

from transformers import EarlyStoppingCallback

training_args = TrainingArguments(
    output_dir="./results",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=64,
    fp16=True,
    num_train_epochs=5,
    weight_decay=0.01,
    logging_steps=50,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    save_total_limit=2,
    seed=42,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
)

trainer.train()

def predict(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=64).to(model.device)
    with torch.no_grad():
        logits = model(**inputs).logits
    pred = torch.argmax(logits, dim=-1).item()
    return "긍정" if pred == 1 else "부정"

print(predict("연출도 좋고 배우들 연기도 훌륭했어요"))
print(predict("완전 시간 낭비였습니다"))
print(predict("정말 재미 없었어요"))

from huggingface_hub import login

login(token="hf_xxxxxxxxxxxxxxxxxxxx")

repo_id = "blueholez/klue-bert-sentiment-blueholez-demo"

model.push_to_hub(repo_id)
tokenizer.push_to_hub(repo_id)

repo_id = "blueholez/klue-bert-sentiment-blueholez-demo"

# 업로드된 모델을 다시 불러와서 검증
reloaded = AutoModelForSequenceClassification.from_pretrained(repo_id)
reloaded_tokenizer = AutoTokenizer.from_pretrained(repo_id)
print("업로드 및 재로드 성공:", repo_id)

import os
import types
import torch
import torch.nn as nn
import pandas as pd
import urllib.request
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel
from trl.trainer.sft_trainer import SFTTrainer
from trl.trainer.sft_config import SFTConfig

model_name = "LGAI-EXAONE/EXAONE-4.0-1.2B"
device = "cuda" if torch.cuda.is_available() else "cpu"

urllib.request.urlretrieve(
    "https://raw.githubusercontent.com/e9t/nsmc/master/ratings_train.txt",
    filename="ratings_train.txt",
)

train_df = pd.read_csv("ratings_train.txt", sep="\t").dropna(subset=["document"])
train_df = train_df[train_df["document"].str.strip().str.len() > 0]

train_df = train_df.sample(n=1000, random_state=42).reset_index(drop=True)

label_to_text = {0: "부정적인 리뷰입니다.", 1: "긍정적인 리뷰입니다."}
train_df["output"] = train_df["label"].map(label_to_text)

def format_prompt(example):
    return {
        "text": (
            f"### 지시문:\n다음 영화 리뷰의 감성을 분석해줘\n\n"
            f"### 입력:\n{example['document']}\n\n"
            f"### 응답:\n{example['output']}"
        )
    }

train_dataset = Dataset.from_pandas(train_df[["document", "output"]]).map(format_prompt)

def patch_input_embeddings(model):
    def find_embedding_module(m):
        for name, module in m.named_modules():
            if isinstance(module, nn.Embedding):
                return name, module
        raise ValueError("Embedding 모듈을 찾을 수 없습니다.")

    embed_name, embed_module = find_embedding_module(model)

    def _get_input_embeddings(self):
        return embed_module

    def _set_input_embeddings(self, new_embeddings):
        parts = embed_name.split(".")
        obj = model
        for p in parts[:-1]:
            obj = getattr(obj, p)
        setattr(obj, parts[-1], new_embeddings)

    model.get_input_embeddings = types.MethodType(_get_input_embeddings, model)
    model.set_input_embeddings = types.MethodType(_set_input_embeddings, model)

    if hasattr(model, "model"):
        model.model.get_input_embeddings = types.MethodType(_get_input_embeddings, model.model)
        model.model.set_input_embeddings = types.MethodType(_set_input_embeddings, model.model)

    return model

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)

base_model = prepare_model_for_kbit_training(base_model)
base_model = patch_input_embeddings(base_model)

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(base_model, lora_config)
model.print_trainable_parameters()

sft_config = SFTConfig(
    output_dir="./lora-results",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    num_train_epochs=3,
    learning_rate=2e-4,
    logging_steps=1,
    save_strategy="epoch",
    report_to="none",
    dataset_text_field="text",
    max_length=512,
)

trainer = SFTTrainer(
    model=model,
    args=sft_config,
    train_dataset=train_dataset,
    processing_class=tokenizer,
)

trainer.train()

adapter_dir = "./lora-adapter"
trainer.save_model(adapter_dir)
tokenizer.save_pretrained(adapter_dir)

def test_sentiment_instruction(review_text):
    prompt = (
        f"### 지시문:\n다음 영화 리뷰의 감성을 분석해줘\n\n"
        f"### 입력:\n{review_text}\n\n"
        f"### 응답:\n"
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=30,
            eos_token_id=tokenizer.eos_token_id,
        )
    generated = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return generated[len(prompt):].strip()

merged_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="cpu",
    trust_remote_code=True,
)
merged_model = patch_input_embeddings(merged_model)

merged_model = PeftModel.from_pretrained(merged_model, adapter_dir)
merged_model = merged_model.merge_and_unload()

merged_dir = "./merged-model"
merged_model.save_pretrained(merged_dir, safe_serialization=True)
tokenizer.save_pretrained(merged_dir)

from huggingface_hub import HfApi, create_repo

hf_username = "blueholez"
merged_repo_id = f"{hf_username}/exaone-nsmc-lora-merged-bluehoez"

create_repo(merged_repo_id, repo_type="model", exist_ok=True)
merged_model.push_to_hub(merged_repo_id)
tokenizer.push_to_hub(merged_repo_id)

import os
print(os.getcwd())                              # 현재 작업 디렉토리
print(os.path.exists("./merged-model-f16.gguf")) # 파일이 실제로 거기 있는지

from huggingface_hub import login, HfApi, create_repo

# 1) 로그인 (토큰 필요 — huggingface.co/settings/tokens 에서 "Write" 권한으로 발급)
#login(token="hf_xxxxxxxxxxxxxxxxxxxx")

# 2) 업로드할 새 저장소 생성
hf_username = "blueholez"   # 본인 계정명으로 변경
repo_id = f"{hf_username}/exaone-nsmc-lora-blueholez-GGUF"

create_repo(repo_id, repo_type="model", exist_ok=True)

# 3) GGUF 파일 업로드
api = HfApi()
api.upload_file(
    path_or_fileobj="./merged-model-f16.gguf",   # llama.cpp 폴더 기준 상대경로
    path_in_repo="exaone-nsmc-lora-f16.gguf",
    repo_id=repo_id,
    repo_type="model",
)

print(f"업로드 완료: https://huggingface.co/{repo_id}")