import gradio as gr
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Define the model class
class GPTLanguageModel(GPT2LMHeadModel):
    def __init__(self, config):
        super().__init__(config)

# Load tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")  # Use your tokenizer path
model = GPTLanguageModel.from_pretrained("gpt2")  # Load the architecture
model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')))  # Load the weights
model.eval()  # Set to evaluation mode

# Define a custom text generation pipeline
class CustomTextGenerationPipeline:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def __call__(self, prompt, max_length=100):
        input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
        generated_ids = self.model.generate(input_ids, max_length=max_length)
        return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)

# Create the pipeline
pipeline = CustomTextGenerationPipeline(model, tokenizer)

# Define the Gradio response function
def respond(message):
    return pipeline(message, max_length=100)

# Create the Gradio interface
demo = gr.Interface(
    fn=respond,
    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt..."),
    outputs="text",
)

if __name__ == "__main__":
    demo.launch()