Upload app.py

app.py

@@ -1,19 +1,19 @@
 import gradio as gr
 import torch
-from transformers import GPT2Tokenizer
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
-# Assuming 'GPTLanguageModel' is already defined
-class GPTLanguageModel(torch.nn.Module):
-    def forward(self, input_ids):
-        # Placeholder forward function
-        pass
+# Define the model class
+class GPTLanguageModel(GPT2LMHeadModel):
+    def __init__(self, config):
+        super().__init__(config)
 
-    def generate(self, input_ids, max_length=100):
-        # This is a placeholder. Replace this with your actual text generation logic.
-        # Right now it just returns the input back, but in your real model, this would
-        # generate new tokens.
-        return input_ids  # Just returning the input as is, to mimic generation
+# 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
@@ -21,17 +21,9 @@ class CustomTextGenerationPipeline:
 
     def __call__(self, prompt, max_length=100):
         input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
-        # Generate text using the model (this is currently simplified)
         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", map_location=torch.device('cpu')))  # Load weights onto CPU
-model.eval()
-
-tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-
 # Create the pipeline
 pipeline = CustomTextGenerationPipeline(model, tokenizer)