Upload 2 files

app.py

@@ -56,6 +56,6 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()
 else:
-    app = demo.launch(share=False)
+    app = demo.launch()

model_utils.py

@@ -110,12 +110,16 @@ class GPT(nn.Module):
 
 def load_model(model_path):
     """Load the trained model"""
-    checkpoint = torch.load(model_path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
-    config = checkpoint['config']
-    model = GPT(config)
-    model.load_state_dict(checkpoint['model_state_dict'])
-    model.eval()
-    return model
+    try:
+        checkpoint = torch.load(model_path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'), weights_only=True)
+        config = GPTConfig(**checkpoint['config'])
+        model = GPT(config)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        model.eval()
+        return model
+    except AttributeError as e:
+        print(f"Error loading model: {e}")
+        return None
 
 def generate_text(model, prompt, max_new_tokens=50, temperature=0.8, top_k=40):
     """Generate text based on a prompt