Update app.py

app.py

@@ -1,398 +1,159 @@
 import gradio as gr
 import torch
-import logging
-import yaml
-import os
-import json
-import jwt
-import redis
-import sqlite3
-from datetime import datetime, timedelta
-from pathlib import Path
 from transformers import AutoTokenizer, T5ForConditionalGeneration
 import networkx as nx
 import matplotlib.pyplot as plt
 import numpy as np
 from PIL import Image
 import io
-import traceback
-from typing import Tuple, Optional, Dict, List, Union
-import colorama
-from colorama import Fore, Style
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-import plotly.graph_objects as go
-import hashlib
-import asyncio
-import aiohttp
-from fastapi import FastAPI, HTTPException, Depends, status
-from fastapi.security import OAuth2PasswordBearer
-from pydantic import BaseModel, EmailStr
-import uvicorn
+from sklearn.feature_extraction.text import TfidfVectorizer
+from scipy.spatial import distance
 
-# Advanced Configuration Models
-@dataclass
-class ModelConfig:
-    name: str
-    max_length: int
-    num_beams: int
-    temperature: float
-    top_k: int
-    top_p: float
-
-@dataclass
-class StyleConfig:
-    node_color: str
-    edge_color: str
-    node_size: int
-    font_size: int
-    layout: str
-
-@dataclass
-class OutputConfig:
-    width: int
-    height: int
-    dpi: int
-    format: str
-    quality: int
-
-# Abstract Base Classes for Extensibility
-class DiagramStrategy(ABC):
-    @abstractmethod
-    def create_diagram(self, components: List[str], style: StyleConfig) -> Image.Image:
-        pass
-
-class NetworkDiagram(DiagramStrategy):
-    def create_diagram(self, components: List[str], style: StyleConfig) -> Image.Image:
-        G = nx.DiGraph()
-        for i in range(len(components)-1):
-            G.add_edge(components[i], components[i+1])
-
-        plt.figure(figsize=(12, 8))
+class DiagramGenerator:
+    def __init__(self):
+        # Initialize device
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         
-        if style.layout == "spring":
-            pos = nx.spring_layout(G)
-        elif style.layout == "circular":
-            pos = nx.circular_layout(G)
-        else:
-            pos = nx.kamada_kawai_layout(G)
-
-        nx.draw_networkx_nodes(G, pos,
-                             node_color=style.node_color,
-                             node_size=style.node_size)
-        nx.draw_networkx_edges(G, pos,
-                             edge_color=style.edge_color,
-                             arrows=True)
-        nx.draw_networkx_labels(G, pos,
-                              font_size=style.font_size)
-
-        buf = io.BytesIO()
-        plt.savefig(buf, format='png', dpi=300)
-        plt.close()
-        buf.seek(0)
-        return Image.open(buf)
-
-class PlotlyDiagram(DiagramStrategy):
-    def create_diagram(self, components: List[str], style: StyleConfig) -> Image.Image:
-        G = nx.DiGraph()
-        for i in range(len(components)-1):
-            G.add_edge(components[i], components[i+1])
-
-        pos = nx.spring_layout(G)
+        # Load model
+        self.model_name = "t5-small"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = T5ForConditionalGeneration.from_pretrained(self.model_name).to(self.device)
         
-        edge_x = []
-        edge_y = []
-        for edge in G.edges():
-            x0, y0 = pos[edge[0]]
-            x1, y1 = pos[edge[1]]
-            edge_x.extend([x0, x1, None])
-            edge_y.extend([y0, y1, None])
-
-        node_x = [pos[node][0] for node in G.nodes()]
-        node_y = [pos[node][1] for node in G.nodes()]
-
-        fig = go.Figure()
-        fig.add_trace(go.Scatter(x=edge_x, y=edge_y,
-                                line=dict(width=0.5, color=style.edge_color),
-                                hoverinfo='none',
-                                mode='lines'))
+        # Initialize vectorizer
+        self.vectorizer = TfidfVectorizer(stop_words='english')
         
-        fig.add_trace(go.Scatter(x=node_x, y=node_y,
-                                mode='markers+text',
-                                marker=dict(size=style.node_size/100,
-                                          color=style.node_color),
-                                text=list(G.nodes()),
-                                textposition="bottom center"))
-
-        fig.update_layout(showlegend=False,
-                         hovermode='closest',
-                         margin=dict(b=0,l=0,r=0,t=0))
-
-        img_bytes = fig.to_image(format="png")  # Requires kaleido
-        return Image.open(io.BytesIO(img_bytes))
+        # Style configurations
+        self.styles = {
+            "flowchart": {
+                "node_color": "lightblue",
+                "edge_color": "gray",
+                "node_size": 3000
+            },
+            "mindmap": {
+                "node_color": "lightgreen",
+                "edge_color": "darkgreen",
+                "node_size": 2500
+            },
+            "sequence": {
+                "node_color": "lightyellow",
+                "edge_color": "orange",
+                "node_size": 3500
+            },
+            "kga": {
+                "node_color": "lightcoral",
+                "edge_color": "darkred",
+                "node_size": 3000
+            }
+        }
 
-# Database Manager
-class DatabaseManager:
-    def __init__(self, db_path: str = "diagrams.db"):
-        self.conn = sqlite3.connect(db_path)
-        self.create_tables()
+    def extract_components(self, text: str) -> list:
+        """Extract components from text using T5 model."""
+        inputs = self.tokenizer(
+            text,
+            max_length=512,
+            truncation=True,
+            return_tensors="pt"
+        ).to(self.device)
 
-    def create_tables(self):
-        cursor = self.conn.cursor()
-        cursor.execute('''
-        CREATE TABLE IF NOT EXISTS users (
-            id INTEGER PRIMARY KEY,
-            username TEXT UNIQUE,
-            email TEXT UNIQUE,
-            password_hash TEXT,
-            created_at TIMESTAMP
+        outputs = self.model.generate(
+            inputs['input_ids'],
+            num_beams=4,
+            max_length=512
         )
-        ''')
-        
-        cursor.execute('''
-        CREATE TABLE IF NOT EXISTS diagrams (
-            id INTEGER PRIMARY KEY,
-            user_id INTEGER,
-            title TEXT,
-            description TEXT,
-            created_at TIMESTAMP,
-            image_path TEXT,
-            FOREIGN KEY (user_id) REFERENCES users (id)
-        )
-        ''')
-        self.conn.commit()
-
-# Cache Manager using Redis
-class CacheManager:
-    def __init__(self, redis_url: str = "redis://localhost"):
-        self.redis_client = redis.from_url(redis_url)
 
-    def get_cached_diagram(self, key: str) -> Optional[bytes]:
-        return self.redis_client.get(key)
-
-    def cache_diagram(self, key: str, diagram: bytes, expire: int = 3600):
-        self.redis_client.set(key, diagram, ex=expire)
-
-# Advanced Diagram Generator
-class AdvancedDiagramGenerator:
-    def __init__(self, config_path: str = "config.yaml"):
-        # Initialize logging first
-        self.setup_logging()
-        
-        # Load configuration
-        self.load_config(config_path)
-        
-        # Setup components (tokenizer, model, etc.)
-        self.setup_components()
-        
-        # Initialize diagram strategies
-        self.strategies = {
-            "network": NetworkDiagram(),
-            "plotly": PlotlyDiagram()
-        }
+        decoded_output = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return [comp.strip() for comp in decoded_output.split(",")]
 
-    def load_config(self, config_path: str):
+    def create_diagram(self, text: str, style: str = "flowchart"):
+        """Create diagram from text with specified style."""
         try:
-            if not os.path.exists(config_path):
-                self.logger.warning(f"Config file not found at {config_path}. Using default configuration.")
-                # Default configuration
-                config_data = {
-                    'model': {
-                        'name': 't5-small',
-                        'max_length': 512,
-                        'num_beams': 4,
-                        'temperature': 1.0,
-                        'top_k': 50,
-                        'top_p': 0.9
-                    },
-                    'styles': {
-                        'network': {
-                            'node_color': '#1f77b4',
-                            'edge_color': '#7f7f7f',
-                            'node_size': 3000,
-                            'font_size': 12,
-                            'layout': 'spring'
-                        }
-                    },
-                    'output': {
-                        'width': 1200,
-                        'height': 800,
-                        'dpi': 300,
-                        'format': 'png',
-                        'quality': 95
-                    }
-                }
+            # Extract components
+            components = self.extract_components(text)
+            if not components:
+                return None, "No components extracted from text."
+
+            # Create figure
+            plt.figure(figsize=(12, 8))
+            G = nx.DiGraph()
+
+            if style == "kga":
+                # Create KGA diagram
+                tfidf_matrix = self.vectorizer.fit_transform(components)
+                similarity_matrix = 1 - distance.squareform(
+                    distance.pdist(tfidf_matrix.toarray(), metric='cosine')
+                )
+                
+                # Add edges based on similarity
+                for i in range(len(components)):
+                    for j in range(i + 1, len(components)):
+                        if similarity_matrix[i][j] > 0.5:
+                            G.add_edge(components[i], components[j])
+                            G.add_edge(components[j], components[i])
             else:
-                with open(config_path) as f:
-                    config_data = yaml.safe_load(f)
+                # Create sequential diagram
+                for i in range(len(components)-1):
+                    G.add_edge(components[i], components[i+1])
 
-            # Ensure all required sections exist with defaults
-            if 'model' not in config_data:
-                config_data['model'] = {}
-            
-            # Set default values for model configuration
-            model_config_data = config_data['model']
-            model_config_data.setdefault('name', 't5-small')
-            model_config_data.setdefault('max_length', 512)
-            model_config_data.setdefault('num_beams', 4)
-            model_config_data.setdefault('temperature', 1.0)
-            model_config_data.setdefault('top_k', 50)
-            model_config_data.setdefault('top_p', 0.9)
+            # Draw diagram
+            pos = nx.spring_layout(G)
+            style_config = self.styles[style]
             
-            # Create ModelConfig instance
-            self.model_config = ModelConfig(**model_config_data)
-
-            # Handle styles configuration
-            if 'styles' not in config_data:
-                config_data['styles'] = {
-                    'network': {
-                        'node_color': '#1f77b4',
-                        'edge_color': '#7f7f7f',
-                        'node_size': 3000,
-                        'font_size': 12,
-                        'layout': 'spring'
-                    }
-                }
+            nx.draw_networkx_nodes(
+                G, pos,
+                node_color=style_config['node_color'],
+                node_size=style_config['node_size']
+            )
             
-            # Create StyleConfig instances
-            self.style_configs = {}
-            for style_name, style_data in config_data['styles'].items():
-                style_data.setdefault('node_color', '#1f77b4')
-                style_data.setdefault('edge_color', '#7f7f7f')
-                style_data.setdefault('node_size', 3000)
-                style_data.setdefault('font_size', 12)
-                style_data.setdefault('layout', 'spring')
-                self.style_configs[style_name] = StyleConfig(**style_data)
-
-            # Handle output configuration
-            if 'output' not in config_data:
-                config_data['output'] = {}
+            nx.draw_networkx_edges(
+                G, pos,
+                edge_color=style_config['edge_color'],
+                arrows=True if style != "kga" else False
+            )
             
-            output_config_data = config_data['output']
-            output_config_data.setdefault('width', 1200)
-            output_config_data.setdefault('height', 800)
-            output_config_data.setdefault('dpi', 300)
-            output_config_data.setdefault('format', 'png')
-            output_config_data.setdefault('quality', 95)
+            nx.draw_networkx_labels(G, pos)
+            plt.title(f"{style.capitalize()} Diagram")
+            plt.axis('off')
+
+            # Save to buffer
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png', bbox_inches='tight', dpi=100)
+            plt.close()
+            buf.seek(0)
             
-            # Create OutputConfig instance
-            self.output_config = OutputConfig(**output_config_data)
-
-            self.config = config_data
-            self.logger.info("Configuration loaded successfully")
-
-        except Exception as e:
-            self.logger.error(f"Error loading configuration: {str(e)}")
-            raise RuntimeError(f"Failed to load configuration: {str(e)}")
-
-    def setup_components(self):
-        # Initialize tokenizer and model
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_config.name)
-        self.model = T5ForConditionalGeneration.from_pretrained(self.model_config.name)
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(self.device)
-
-    def setup_logging(self):
-        logging.basicConfig(level=logging.INFO)
-        self.logger = logging.getLogger(__name__)
-
-    async def extract_components(self, text: str) -> List[str]:
-        inputs = self.tokenizer(
-            f"convert to diagram: {text}",
-            return_tensors="pt",
-            max_length=self.model_config.max_length,
-            truncation=True
-        ).to(self.device)
-
-        with torch.no_grad():
-            outputs = self.model.generate(
-                inputs.input_ids,
-                max_length=150,
-                num_beams=self.model_config.num_beams,
-                temperature=self.model_config.temperature,
-                top_k=self.model_config.top_k,
-                top_p=self.model_config.top_p
-            )
-
-        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        components = [comp.strip() for comp in decoded.replace('->', ',').split(',')]
-        return [comp for comp in components if comp]
+            return Image.open(buf), "Diagram generated successfully!"
 
-    def save_to_database(self, user_id: int, description: str, diagram: Image.Image):
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        image_path = f"diagrams/user_{user_id}/{timestamp}.png"
-        os.makedirs(os.path.dirname(image_path), exist_ok=True)
-        diagram.save(image_path)
-
-        cursor = self.db.conn.cursor()
-        cursor.execute('''
-        INSERT INTO diagrams (user_id, description, created_at, image_path)
-        VALUES (?, ?, ?, ?)
-        ''', (user_id, description, datetime.now(), image_path))
-        self.db.conn.commit()
-
-    async def generate_diagram(self, text: str, style: str, strategy: str) -> Tuple[Optional[Image.Image], str]:
-        try:
-            components = await self.extract_components(text)
-            diagram = self.strategies[strategy].create_diagram(components, self.style_configs[style])
-            return diagram, "Diagram generated successfully!"
         except Exception as e:
-            self.logger.error(f"Error generating diagram: {str(e)}")
-            return None, f"Error: {str(e)}"
+            return None, f"Error generating diagram: {str(e)}"
 
-# FastAPI Integration
-app = FastAPI()
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
-
-class DiagramRequest(BaseModel):
-    text: str
-    style: str = "network"
-    strategy: str = "network"
-
-# Initialize the generator
-generator = AdvancedDiagramGenerator()
-
-# Gradio Interface
 def create_gradio_interface():
-    def generate(text: str, style: str, strategy: str) -> Tuple[Optional[Image.Image], str]:
-        return asyncio.run(generator.generate_diagram(text, style, strategy))
-
+    generator = DiagramGenerator()
+    
     iface = gr.Interface(
-        fn=generate,
+        fn=generator.create_diagram,
         inputs=[
             gr.Textbox(
                 label="Enter your diagram description",
-                placeholder="e.g., 'Create a flowchart for software development lifecycle'",
+                placeholder="e.g., 'Create a knowledge graph for artificial intelligence concepts'",
                 lines=3
             ),
             gr.Dropdown(
-                choices=list(generator.style_configs.keys()),
+                choices=list(generator.styles.keys()),
                 label="Diagram Style",
-                value="network"
-            ),
-            gr.Dropdown(
-                choices=list(generator.strategies.keys()),
-                label="Visualization Strategy",
-                value="network"
+                value="flowchart"
             )
         ],
         outputs=[
             gr.Image(label="Generated Diagram", type="pil"),
             gr.Textbox(label="Status")
         ],
-        title="Advanced Enterprise Diagram Generator",
+        title="AI-Powered Diagram Generator",
         description="""
-        Enterprise-grade diagram generation tool with advanced features:
-        - Multiple visualization strategies
-        - Caching system
-        - Database storage
-        - Multiple output formats
-        - Custom styling options
-        """,
-        theme=gr.themes.Glass()
+        Create various types of diagrams from text descriptions.
+        Supports flowcharts, mindmaps, sequence diagrams, and knowledge graphs.
+        """
     )
     return iface
 
 if __name__ == "__main__":
     iface = create_gradio_interface()
-    iface.launch(share=True)
+    iface.launch()