import torch
from tiny_model import TinyLanguageModel, generate_text

def load_model(model_path):
    """بارگذاری مدل آموزش دیده"""
    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
    
    model = TinyLanguageModel(
        checkpoint['vocab_size'],
        embedding_dim=128,
        hidden_dim=256,
        num_layers=2
    )
    
    model.load_state_dict(checkpoint['model_state_dict'])
    return model, checkpoint['char_to_idx'], checkpoint['idx_to_char']

def interactive_generation():
    """تعامل با مدل برای تولید متن"""
    print("در حال بارگذاری مدل...")
    model, char_to_idx, idx_to_char = load_model('tiny_language_model.pth')
    
    print("مدل بارگذاری شد! برای خروج 'exit' را وارد کنید.")
    
    while True:
        prompt = input("\nمتن شروع را وارد کنید: ")
        if prompt.lower() == 'exit':
            break
        
        length = int(input("طول متن تولیدی: "))
        temperature = float(input("دمای تولید (0.1-2.0): "))
        
        generated_text = generate_text(
            model, prompt, char_to_idx, idx_to_char, 
            length=length, temperature=temperature
        )
        
        print(f"\nمتن تولید شده: {generated_text}")

if __name__ == "__main__":
    interactive_generation()