Create app.py

app.py

@@ -0,0 +1,72 @@
+import gradio as gr
+from transformers import PreTrainedTokenizerFast, AutoConfig
+from safetensors.torch import load_model  # Import load_model for safetensors
+import torch
+from EnhancedTransformerModel import EnhancedTransformerModel  # Custom model class
+
+# --- Load Model and Tokenizer ---
+MODEL_PATH = "model.safetensors"  # Path to the safetensors model file
+CONFIG_PATH = "config.json"       # Path to the model configuration
+TOKENIZER_PATH = "tokenizer"      # Path to the tokenizer directory
+
+# Load configuration
+config = AutoConfig.from_pretrained(CONFIG_PATH)
+
+# Load tokenizer
+tokenizer = PreTrainedTokenizerFast.from_pretrained(TOKENIZER_PATH)
+
+# Initialize the custom model
+model = EnhancedTransformerModel(
+    vocab_size=config.vocab_size,
+    max_seq_length=config.max_position_embeddings,
+    d_model=config.hidden_size,
+    num_heads=config.num_attention_heads,
+    num_layers=config.num_hidden_layers,
+    ff_dim=config.intermediate_size,
+    dropout=0.1
+)
+
+# Load the model weights from safetensors
+state_dict = load_model(MODEL_PATH)
+model.load_state_dict(state_dict)
+model.eval()
+model.to("cuda" if torch.cuda.is_available() else "cpu")
+
+# --- Inference Function ---
+def generate_text(prompt, max_length=50):
+    """
+    Generate text based on the input prompt using the trained model.
+    """
+    # Tokenize the input prompt
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=384)
+    input_ids = inputs["input_ids"].to(model.device)
+    attention_mask = inputs["attention_mask"].to(model.device)
+
+    # Generate output tokens
+    with torch.no_grad():
+        outputs = model(input_ids, attention_mask)
+        logits = outputs[:, -1, :]  # Get the logits for the last token
+        next_token = torch.argmax(logits, dim=-1)  # Greedy decoding
+
+    # Decode the generated token
+    generated_text = tokenizer.decode(next_token, skip_special_tokens=True)
+    return generated_text
+
+# --- Gradio Interface ---
+with gr.Blocks() as demo:
+    gr.Markdown("# Snowflake-G0-stable Language Model")
+    gr.Markdown("This is an enhanced transformer language model trained on the DialogMLM-50K dataset. Try it out below!")
+
+    with gr.Row():
+        input_prompt = gr.Textbox(label="Input Prompt", placeholder="Enter your text here...")
+        output_text = gr.Textbox(label="Generated Text")
+
+    submit_button = gr.Button("Generate")
+
+    def on_submit(prompt):
+        return generate_text(prompt)
+
+    submit_button.click(on_submit, inputs=input_prompt, outputs=output_text)
+
+# Launch the app
+demo.launch()