Create app.py

app.py

@@ -0,0 +1,49 @@
+# Import required libraries
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+
+# Load the Hugging Face model and tokenizer
+model_name = "TroyDoesAI/Mermaid-Llama-3-8B"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Define the pipeline for text generation
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+# Function to generate Mermaid.js code and return diagram
+def generate_mermaid_code(user_input):
+    # Generate Mermaid.js code using the model
+    prompt = f"Create a Mermaid.js diagram based on the following story:\n{user_input}\nMermaid.js Code:"
+    generated = pipe(prompt, max_length=200, num_return_sequences=1, do_sample=True)
+    
+    # Extract the generated Mermaid.js code
+    mermaid_code = generated[0]['generated_text'].split("Mermaid.js Code:")[-1].strip()
+    
+    # Create the full output for Gradio with Mermaid.js rendering
+    mermaid_output = f"""
+    ```mermaid
+    {mermaid_code}
+    ```
+    """
+    return mermaid_code, mermaid_output
+
+# Define the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Mermaid.js Diagram Generator")
+    gr.Markdown("Provide a story, and the AI will generate a Mermaid.js diagram code.")
+    
+    with gr.Row():
+        input_text = gr.Textbox(label="Story Input", lines=5, placeholder="Enter your story here...")
+        output_code = gr.Textbox(label="Generated Mermaid.js Code", lines=5)
+    
+    output_diagram = gr.Markdown(label="Mermaid.js Diagram")
+    
+    # Button to trigger generation
+    generate_button = gr.Button("Generate Diagram")
+    
+    # Link input to output
+    generate_button.click(generate_mermaid_code, inputs=input_text, outputs=[output_code, output_diagram])
+
+# Launch the Gradio interface
+demo.launch()