import gradio as gr
import torch
from transformers import GPT2Tokenizer

# Assuming 'GPTLanguageModel' is already defined
class GPTLanguageModel(torch.nn.Module):
    def forward(self, input_ids):
        pass

    def generate(self, input_ids, max_length=100):
        return torch.tensor([[input_ids]])  # This is a placeholder for generation

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)

# Load model and tokenizer
model = GPTLanguageModel()
model.load_state_dict(torch.load("model.pth"))  # Load the weights
model.eval()

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

# 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()