import os
from transformers import GPT2Tokenizer, GPT2LMHeadModel

# Define paths
model_dir = "./fine-tuned-tiny-gpt2"

# Load the fine-tuned model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained(model_dir)
model = GPT2LMHeadModel.from_pretrained(model_dir)

# Optionally add a pad token if not present
if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({'pad_token': '<PAD>'})
    model.resize_token_embeddings(len(tokenizer))

def generate_response(prompt, model, tokenizer, max_length=512, max_new_tokens=50, temperature=0.1, top_k=50, top_p=0.9):
    inputs = tokenizer(prompt, return_tensors='pt', truncation=True, padding='max_length', max_length=max_length)
    
    outputs = model.generate(
        input_ids=inputs['input_ids'],
        attention_mask=inputs['attention_mask'],
        max_length=max_length + max_new_tokens,
        num_return_sequences=1,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        do_sample=True
    )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

def chat_with_model(model, tokenizer):
    print("Chat with the model (type 'exit' to quit):")
    
    while True:
        user_input = input("You: ")
        
        if user_input.lower() == 'exit':
            break
        
        response = generate_response(user_input, model, tokenizer)
        print("Model: ", response)
        
# Start the chat
chat_with_model(model, tokenizer)