Upload app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import torch
 from transformers import GPT2Tokenizer
 
+# Define your model
 class GPTLanguageModel(torch.nn.Module):
     def __init__(self, vocab_size, hidden_size):
         super(GPTLanguageModel, self).__init__()
@@ -15,14 +16,25 @@ class GPTLanguageModel(torch.nn.Module):
         # Custom generation logic here
         return input_ids  # Placeholder
 
+# Define the 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)
+
 # Load tokenizer and model
 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
+model = GPTLanguageModel(vocab_size=vocab_size, hidden_size=768)  # Set sizes appropriately
 
 # Load model weights
 try:
-    model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False)
+    model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu'), weights_only=True), strict=False)
 except RuntimeError as e:
     print(f"Error loading model weights: {e}")