import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel, AutoModelForSeq2SeqLM, AutoTokenizer, GPT2Config
import torch
from safetensors.torch import load_file as safetensors_load_file  # Import safetensors loading function
from langchain.memory import ConversationBufferMemory

# Move model to device (GPU if available)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# Load the tokenizer (use pre-trained tokenizer for GPT-2 family)
tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")

# Load the configuration for the model (DistilGPT2 is a smaller GPT-2)
config = GPT2Config.from_pretrained("distilgpt2")

# Initialize the model using the configuration
model = GPT2LMHeadModel(config)

# Load the weights from the safetensors file
model_path = "./model.safetensors"  # Path to your local model file
state_dict = safetensors_load_file(model_path)  # Use safetensors loader
model.load_state_dict(state_dict)  # Load the state dict into the model

# Move model to the device (GPU or CPU)
model.to(device)

# Set up conversational memory using LangChain's ConversationBufferMemory
memory = ConversationBufferMemory()

# Define the chatbot function with memory
def chat_with_distilgpt2(input_text):
    # Retrieve conversation history
    conversation_history = memory.load_memory_variables({})['history']
    
    # Combine the (possibly summarized) history with the current user input
    full_input = f"{conversation_history}\nUser: {input_text}\nAssistant:"
    
    # Tokenize the input and convert to tensor
    input_ids = tokenizer.encode(full_input, return_tensors="pt").to(device)
    
    # Generate the response using the model with adjusted parameters
    outputs = model.generate(
        input_ids,
        max_length=input_ids.shape[1] + 100,  # Limit total length
        max_new_tokens=100,
        num_return_sequences=1,
        no_repeat_ngram_size=3,
        repetition_penalty=1.2,
        early_stopping=True,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id
    )
    
    # Decode the model output
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Update the memory with the user input and model response
    memory.save_context({"input": input_text}, {"output": response})
    
    return response

# Set up the Gradio interface
interface = gr.Interface(
    fn=chat_with_distilgpt2,
    inputs=gr.Textbox(label="Chat with DistilGPT-2"),
    outputs=gr.Textbox(label="DistilGPT-2's Response"),
    title="DistilGPT-2 Chatbot with Memory",
    description="This is a simple chatbot powered by the DistilGPT-2 model with conversational memory, using LangChain.",
)

# Launch the Gradio app
interface.launch()