import torch
from transformers import T5Tokenizer
from peft import PeftModel, PeftConfig

# Check for MPS availability (Apple Silicon)
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
#CUDA
#device = torch.device("CUDA")

# Path to your saved model
model_path = "./question-paraphrase-T5"

tokenizer = T5Tokenizer.from_pretrained("t5-base")

config = PeftConfig.from_pretrained(model_path)
from transformers import AutoModelForSeq2SeqLM
model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path)
model = PeftModel.from_pretrained(model, model_path)
model = model.to(device)
model.eval()  # Set to evaluation mode


def generate_paraphrase(text):
    """Generate a paraphrase for the given text."""
    input_text = f"paraphrase : {text}"
    print(f"Input: {input_text}")
    
    inputs = tokenizer(input_text, return_tensors="pt").to(device)
    
    # Use keyword arguments for generate
    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs.input_ids,
            attention_mask=inputs.attention_mask,
            max_length=128,
            num_beams=4,  # Use beam search for better quality
            no_repeat_ngram_size=2,  # Avoid repeating phrases
            temperature=0.7  # Add some randomness for diversity
        )
    
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result


# Interactive mode
print("\nInteractive Mode: Enter questions to paraphrase (type 'exit' to quit)")
while True:
    user_input = input("\nEnter text to paraphrase: ")
    if user_input.lower() == 'exit':
        break
    
    paraphrase = generate_paraphrase(user_input)
    print(f"Paraphrase: {paraphrase}")