T5

# T5 모델 활용 예제
from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch

# 모델 및 토크나이저 로드 (t5-small, t5-base, t5-large, t5-3b, t5-11b)
model_name = "t5-base"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

def generate_text(input_text: str, max_length: int = 128) -> str:
    """T5로 텍스트 생성"""
    input_ids = tokenizer.encode(input_text, return_tensors="pt")

    outputs = model.generate(
        input_ids,
        max_length=max_length,
        num_beams=4,
        early_stopping=True,
        no_repeat_ngram_size=2
    )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# ===== 1. 번역 (Translation) =====
print("=" * 50)
print("1. 번역 (Translation)")
print("=" * 50)

en_to_de = "translate English to German: The house is wonderful."
de_result = generate_text(en_to_de)
print(f"입력: {en_to_de}")
print(f"출력: {de_result}")

en_to_fr = "translate English to French: Machine learning is fascinating."
fr_result = generate_text(en_to_fr)
print(f"\n입력: {en_to_fr}")
print(f"출력: {fr_result}")

# ===== 2. 요약 (Summarization) =====
print("\n" + "=" * 50)
print("2. 요약 (Summarization)")
print("=" * 50)

article = """
summarize: The Amazon rainforest produces about 20% of the world's oxygen
and contains 10% of all species on Earth. Climate change and deforestation
pose serious threats to this vital ecosystem. Scientists warn that losing
the Amazon could accelerate global warming significantly. Conservation
efforts are critical for the planet's future.
"""
summary = generate_text(article, max_length=64)
print(f"원문: {article.strip()}")
print(f"요약: {summary}")

# ===== 3. 문법 교정 (Grammar Correction) =====
print("\n" + "=" * 50)
print("3. 문법 교정 (Grammar Correction)")
print("=" * 50)

# FLAN-T5 또는 파인튜닝된 모델에서 더 좋은 결과
grammar_input = "cola sentence: He don't knows nothing about it."
corrected = generate_text(grammar_input)
print(f"입력: {grammar_input}")
print(f"출력: {corrected}")

# ===== 4. 질문 답변 (Question Answering) =====
print("\n" + "=" * 50)
print("4. 질문 답변 (Question Answering)")
print("=" * 50)

qa_input = """
question: What is the capital of France?
context: France is a country in Western Europe. Its capital city is Paris,
which is known for the Eiffel Tower and the Louvre Museum.
"""
answer = generate_text(qa_input, max_length=32)
print(f"질문 & 컨텍스트: {qa_input.strip()}")
print(f"답변: {answer}")

# ===== 5. FLAN-T5 (향상된 버전) =====
print("\n" + "=" * 50)
print("5. FLAN-T5 (Instruction-tuned)")
print("=" * 50)

# FLAN-T5는 instruction following이 더 좋음
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

flan_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
flan_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")

def flan_generate(prompt: str) -> str:
    inputs = flan_tokenizer(prompt, return_tensors="pt")
    outputs = flan_model.generate(**inputs, max_length=100)
    return flan_tokenizer.decode(outputs[0], skip_special_tokens=True)

instruction = "Explain what machine learning is in simple terms."
response = flan_generate(instruction)
print(f"질문: {instruction}")
print(f"답변: {response}")