Create app.py

app.py

@@ -0,0 +1,311 @@
+import torch
+import gradio as gr
+import matplotlib.pyplot as plt
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainingArguments, Seq2SeqTrainer, DataCollatorForSeq2Seq
+from datasets import Dataset
+import os
+import base64
+import io
+import requests
+#from IPython.display import display, Markdown, HTML  # Remove IPython dependency
+import time
+
+# Check if GPU is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+## Loading the Pre-trained Model
+
+model_name = "facebook/bart-large"  # You could also use "t5-base" or other seq2seq models
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
+
+## Define Training Data (Optional for Fine-tuning)
+
+# Sample training data: [(text_description, mermaid_code), ...]
+training_data = [
+    (
+        "A flowchart showing user login process with success and failure paths",
+        """graph TD
+        A[Start] --> B{User has account?}
+        B -->|Yes| C[Enter credentials]
+        B -->|No| D[Register]
+        C --> E{Valid credentials?}
+        E -->|Yes| F[Login successful]
+        E -->|No| G[Login failed]
+        D --> C
+        """
+    ),
+    (
+        "A sequence diagram showing client-server authentication",
+        """sequenceDiagram
+        participant Client
+        participant Server
+        Client->>Server: Authentication Request
+        Server->>Client: Challenge
+        Client->>Server: Challenge Response
+        Server->>Client: Auth Success/Failure
+        """
+    ),
+    (
+        "A simple entity relationship diagram for a blog system",
+        """erDiagram
+        AUTHOR ||--o{ POST : writes
+        POST ||--o{ COMMENT : contains
+        AUTHOR ||--o{ COMMENT : writes
+        """
+    ),
+    # Add more examples for better fine-tuning
+]
+
+## Fine-tuning (Optional but Recommended)
+
+def fine_tune_model():
+    # Prepare dataset for fine-tuning
+    dataset_dict = {
+        "input_text": [item[0] for item in training_data],
+        "target_text": [item[1] for item in training_data]
+    }
+
+    dataset = Dataset.from_dict(dataset_dict)
+
+    # Tokenize the dataset
+    def preprocess_function(examples):
+        inputs = examples["input_text"]
+        targets = examples["target_text"]
+        model_inputs = tokenizer(inputs, max_length=128, truncation=True, padding="max_length")
+
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(targets, max_length=256, truncation=True, padding="max_length")
+
+        model_inputs["labels"] = labels["input_ids"]
+        return model_inputs
+
+    tokenized_dataset = dataset.map(preprocess_function, batched=True)
+
+    # Define training arguments
+    training_args = Seq2SeqTrainingArguments(
+        output_dir="./results",
+        evaluation_strategy="epoch",
+        learning_rate=5e-5,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        weight_decay=0.01,
+        save_total_limit=3,
+        num_train_epochs=3,
+        predict_with_generate=True,
+        no_cuda=not torch.cuda.is_available() # Added to handle cases when no GPU is available
+    )
+
+    # Define data collator
+    data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+
+    # Create trainer
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset,
+        data_collator=data_collator,
+        tokenizer=tokenizer,
+    )
+
+    # Start fine-tuning
+    trainer.train()
+
+    # Save fine-tuned model
+    model.save_pretrained("./fine_tuned_model")
+    tokenizer.save_pretrained("./fine_tuned_model")
+
+    return model, tokenizer
+
+# Uncomment the line below to run fine-tuning
+# model, tokenizer = fine_tune_model()
+
+## Text to Diagram Function
+
+def get_entity_relationship_diagram():
+    """
+    Return a predefined entity relationship diagram for a blog system
+    """
+    return """erDiagram
+    AUTHOR ||--o{ POST : writes
+    POST ||--o{ COMMENT : contains
+    USER ||--o{ COMMENT : writes
+    USER ||--o{ AUTHOR : can_be
+    POST }|--|| CATEGORY : belongs_to
+    """
+
+def get_flowchart_diagram():
+    """
+    Return a predefined flowchart diagram
+    """
+    return """graph TD
+    A[Start] --> B{User has account?}
+    B -->|Yes| C[Enter credentials]
+    B -->|No| D[Register]
+    C --> E{Valid credentials?}
+    E -->|Yes| F[Login successful]
+    E -->|No| G[Login failed]
+    D --> C
+    """
+
+def get_sequence_diagram():
+    """
+    Return a predefined sequence diagram
+    """
+    return """sequenceDiagram
+    participant User
+    participant System
+    participant Database
+    User->>System: Request data
+    System->>Database: Query data
+    Database->>System: Return results
+    System->>User: Display results
+    """
+
+def text_to_diagram(text_description):
+    """
+    Convert text description to a diagram using pattern matching or model
+    """
+    # For demonstration, use pattern matching for common cases
+    lower_text = text_description.lower()
+
+    # Pattern match common diagram types based on the input text
+    if "entity" in lower_text and "relation" in lower_text and "blog" in lower_text:
+        diagram_code = get_entity_relationship_diagram()
+    elif "flow" in lower_text and "login" in lower_text:
+        diagram_code = get_flowchart_diagram()
+    elif "sequence" in lower_text and "client" in lower_text and "server" in lower_text:
+        diagram_code = get_sequence_diagram()
+    else:
+        # Use the model for other cases
+        try:
+            # Tokenize input text
+            inputs = tokenizer(text_description, return_tensors="pt", max_length=128, truncation=True).to(device)
+
+            # Generate diagram code
+            outputs = model.generate(
+                inputs["input_ids"],
+                max_length=256,
+                num_beams=5,
+                early_stopping=True
+            )
+
+            # Decode the outputs to get the mermaid diagram code
+            diagram_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+            # For non-fine-tuned models, the output is unlikely to be valid Mermaid code
+            # So we'll apply pattern matching to generate appropriate Mermaid code
+            if "flowchart" in lower_text or "flow" in lower_text:
+                diagram_code = """graph TD
+                A[Start] --> B[Process]
+                B --> C[End]
+                """
+            elif "sequence" in lower_text:
+                diagram_code = """sequenceDiagram
+                participant A
+                participant B
+                A->>B: Message
+                B->>A: Response
+                """
+            elif "entity" in lower_text or "er" in lower_text:
+                diagram_code = """erDiagram
+                ENTITY1 ||--o{ ENTITY2 : relates
+                """
+            else:
+                # Default to a simple flowchart
+                diagram_code = """graph TD
+                A[Start] --> B[Process]
+                B --> C[End]
+                """
+        except Exception as e:
+            print(f"Error generating diagram code: {e}")
+            # Fallback to a simple diagram
+            diagram_code = """graph TD
+            A[Error] --> B[Could not generate diagram]
+            """
+
+    # Render the diagram to an image
+    try:
+        # Use Mermaid.ink API to render the diagram
+        img_url = render_mermaid_to_url(diagram_code)
+
+        # Download the image and convert to a data URL for Gradio
+        try:
+            response = requests.get(img_url, timeout=10)
+            if response.status_code == 200:
+                image_data = response.content
+                # Save temporarily to a file that Gradio can display
+                temp_img_path = "temp_diagram.png"  # Fixed filename for simplicity
+                with open(temp_img_path, "wb") as f:
+                    f.write(image_data)
+                return diagram_code, temp_img_path
+            else:
+                return diagram_code, None
+        except Exception as e:
+            print(f"Error downloading image: {e}")
+            return diagram_code, None
+    except Exception as e:
+        print(f"Error rendering diagram: {e}")
+        return diagram_code, None
+
+def render_mermaid_to_url(mermaid_code):
+    """
+    Render mermaid code to an image URL using the Mermaid.live API
+    """
+    try:
+        # Encode the mermaid code to be used in a URL
+        encoded_code = base64.urlsafe_b64encode(mermaid_code.encode()).decode()
+
+        # Generate a URL for the Mermaid.ink service
+        mermaid_url = f"https://mermaid.ink/img/{encoded_code}"
+
+        return mermaid_url
+    except Exception as e:
+        print(f"Error encoding mermaid code: {e}")
+        # Return a fallback URL or None
+        return None
+
+## Gradio Interface
+
+def gradio_interface(text_input):
+    """
+    Process user input and return diagram output via Gradio
+    """
+    try:
+        diagram_code, img_path = text_to_diagram(text_input)
+
+        # Display the diagram code for debugging
+        print("Generated diagram code:")
+        print(diagram_code)
+
+        if img_path:
+            print(f"Image saved to: {img_path}")
+            return diagram_code, img_path
+        else:
+            # If image generation failed, return code only
+            return diagram_code, None
+    except Exception as e:
+        print(f"Error in Gradio interface: {e}")
+        return f"Error generating diagram: {str(e)}", None
+
+# Create the Gradio interface with error handling
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Textbox(lines=5, placeholder="Enter your diagram description here..."),
+    outputs=[
+        gr.Textbox(label="Generated Mermaid Code"),
+        gr.Image(label="Diagram Visualization", type="filepath")
+    ],
+    title="Text to Diagram Converter",
+    description="Convert natural language descriptions to diagrams using AI",
+    examples=[
+        ["A flowchart showing user login process with success and failure paths"],
+        ["A sequence diagram showing client-server authentication"],
+        ["A simple entity relationship diagram for a blog system"]
+    ],
+    allow_flagging="never"
+)
+
+# Launch the interface
+iface.launch()