Upload 2 files

app.py

@@ -1,28 +1,32 @@
 import gradio as gr
 import torch
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from transformers import GPT2Tokenizer
 
-# Define the model class
-class GPTLanguageModel(GPT2LMHeadModel):
-    def __init__(self, config):
-        super().__init__(config)
+class GPTLanguageModel(torch.nn.Module):
+    def __init__(self, vocab_size, hidden_size):
+        super(GPTLanguageModel, self).__init__()
+        self.embedding = torch.nn.Embedding(vocab_size, hidden_size)
+        # Add other layers as needed
+
+    def forward(self, input_ids):
+        return self.embedding(input_ids)  # Placeholder for the forward pass
+
+    def generate(self, input_ids, max_length=100):
+        # Custom generation logic here
+        return input_ids  # Placeholder
 
 # 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)
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")  # Or your custom tokenizer
+vocab_size = tokenizer.vocab_size
+model = GPTLanguageModel(vocab_size=vocab_size, hidden_size=768)  # Set the sizes appropriately
+
+# Load model weights
+try:
+    model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False)
+except RuntimeError as e:
+    print(f"Error loading model weights: {e}")
+
+model.eval()
 
 # Create the pipeline
 pipeline = CustomTextGenerationPipeline(model, tokenizer)

config.json

@@ -0,0 +1,9 @@
+{
+  "block_size": 256,
+  "model_type": "gpt",
+  "n_embd": 384,
+  "n_head": 6,
+  "n_layer": 6,
+  "transformers_version": "4.44.2",
+  "vocab_size": 95
+}