Create app.py

app.py

@@ -0,0 +1,66 @@
+import torch
+import gradio as gr
+from config.model_config import ModelConfig
+from src.data.tokenizer import CharacterTokenizer
+from src.utils.helpers import generate, setup_logging
+
+# Setup logging
+logger = setup_logging()
+
+
+def load_model():
+    config = ModelConfig()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    logger.info(f"Using device: {device}")
+
+    # Load tokenizer
+    with open(config.data_path) as f:
+        text = f.read()
+    tokenizer = CharacterTokenizer(text)
+
+    # Load model
+    try:
+        model = torch.load(config.checkpoint_path, map_location=device)
+        model.eval()
+        return model, tokenizer, device
+    except Exception as e:
+        logger.error(f"Error loading model: {e}")
+        raise
+
+
+def generate_text(prompt, max_tokens=200, temperature=0.8):
+    try:
+        result = generate(model, tokenizer, prompt, max_tokens, device)
+        return prompt + result
+    except Exception as e:
+        logger.error(f"Error during generation: {e}")
+        return f"Error: {str(e)}"
+
+
+# Load model globally
+try:
+    model, tokenizer, device = load_model()
+    logger.info("Model loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to load model: {e}")
+    raise
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(label="Enter your prompt", placeholder="Type your prompt here..."),
+        gr.Slider(minimum=10, maximum=1000, value=200, step=10, label="Max Tokens"),
+    ],
+    outputs=gr.Textbox(label="Generated Text"),
+    title="Shakespeare GPT",
+    description="Enter a prompt and generate text using a custom GPT model",
+    examples=[
+        ["Hello, my name is", 200, 0.8],
+        ["Once upon a time", 500, 0.8],
+        ["The meaning of life is", 300, 0.8],
+    ],
+)
+
+if __name__ == "__main__":
+    demo.launch()