new updates

app.py

@@ -1,92 +1,124 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import json
-import random
 import os
 
-# Initialize the client (uses HUGGING_FACE_HUB_TOKEN from environment)
+# Initialize the client
+# We use a powerful instruction-tuned model. Mistral-7B-Instruct-v0.3 is excellent for reasoning.
+# If that fails, we can fallback to others, but let's start strong.
 token = os.getenv("HUGGING_FACE_HUB_TOKEN")
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=token)
+client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3", token=token)
 
-SYSTEM_PROMPT = """You are ModelForge, an expert AI architecture assistant. Your goal is to analyze machine learning problems and generate detailed, deployable solutions in strict JSON format.
+SYSTEM_PROMPT = """You are ModelForge, an elite AI Solutions Architect with deep expertise in MLOps, research-level machine learning, and software engineering.
+
+Your goal is to analyze the user's problem and design a **complete, production-ready, and theoretically sound** machine learning solution. You must handle complex, multimodal, and research-level requests with ease.
+
+### Output Format
+You must return a SINGLE valid JSON object. Do not include any markdown formatting (like ```json) outside the object.
 
-You must analyze the user's request and return a JSON object with the following structure:
 {
   "analysis": {
     "dataType": "image" | "text" | "tabular" | "audio" | "video" | "time_series" | "multimodal",
-    "taskType": "classification" | "regression" | "nlp" | "vision" | "forecasting" | "multimodal_reasoning",
+    "taskType": "classification" | "regression" | "nlp" | "vision" | "forecasting" | "multimodal_reasoning" | "reinforcement_learning" | "generative",
     "complexity": "low" | "medium" | "high" | "research",
-    "domain": "string (e.g., medical, finance, etc.)"
+    "domain": "string (e.g., 'High-Frequency Trading', 'Autonomous Driving', 'Genomic Research')",
+    "keyChallenges": ["challenge 1", "challenge 2"]
   },
   "recommendations": [
     {
-      "name": "Model Name",
-      "description": "Detailed technical description...",
-      "pros": ["pro1", "pro2", "pro3"],
-      "cons": ["con1", "con2"],
-      "architectureDiagram": "Mermaid graph definition...",
-      "mlopsBestPractices": ["step 1", "step 2", ...],
-      "trainingCode": "Python code snippet..."
+      "name": "Model/Architecture Name",
+      "description": "A comprehensive technical description. For research problems, describe the novel architecture (e.g., 'Dual-Encoder with Cross-Attention Adapters'). For production, specify the exact backbone (e.g., 'ResNet-50v2 with FPN').",
+      "pros": ["Critical advantage 1", "Critical advantage 2", "Critical advantage 3"],
+      "cons": ["Trade-off 1", "Trade-off 2"],
+      "architectureDiagram": "A detailed Mermaid.js graph (graph TD). Use square brackets [] for ALL nodes. NO curly braces {}. Include data ingestion, preprocessing, backbone, heads, and post-processing.",
+      "mlopsBestPractices": [
+        "Data Versioning: Strategy (e.g., DVC/Delta Lake)",
+        "Experiment Tracking: Tools (e.g., MLflow/W&B)",
+        "Deployment: Strategy (e.g., Canary, Blue-Green, Edge)",
+        "Monitoring: Metrics to watch (e.g., Drift, Latency)",
+        "Hardware: Recommended GPU/TPU"
+      ],
+      "trainingCode": "Production-grade Python code snippet (PyTorch/TensorFlow/JAX). Include imports, model definition, and a dummy training loop."
     }
   ]
 }
 
-Provide 2-3 distinct recommendations. For research-level problems, propose novel architectures.
-Ensure the Mermaid diagram uses valid syntax (no curly braces for nodes, use square brackets []).
+### Guidelines for "Robust & Complex"
+1. **Deep Analysis**: Don't just say "tabular". Say "High-dimensional tabular data with potential covariate shift".
+2. **Novel Architectures**: If the user asks for "latent program synthesis", design a "Neural Module Network with Discrete Latent Variables". Do not recommend generic models for research problems.
+3. **Complete Pipelines**: The MLOps section must be actionable and specific to the problem (e.g., "Use ONNX Runtime for <10ms latency").
+4. **Valid JSON**: Your response must be parseable by `json.loads()`.
 """
 
 FEW_SHOT_EXAMPLES = """
-Example Input: "Detect fraud in credit card transactions"
-Example Output:
+User Input: "Create a unified multimodal reasoning model that composes and executes latent programs across vision, text, and audio."
+
+JSON Response:
 {
-  "analysis": { "dataType": "tabular", "taskType": "classification", "complexity": "medium", "domain": "finance" },
+  "analysis": {
+    "dataType": "multimodal",
+    "taskType": "multimodal_reasoning",
+    "complexity": "research",
+    "domain": "Neuro-Symbolic AI",
+    "keyChallenges": ["Cross-modal alignment", "Differentiable program synthesis", "Latent space collapse"]
+  },
   "recommendations": [
     {
-      "name": "XGBoost Fraud Detector",
-      "description": "Gradient boosting ensemble optimized for imbalanced tabular data...",
-      "pros": ["High interpretability", "Handles missing data"],
-      "cons": ["Feature engineering required"],
-      "architectureDiagram": "graph TD\\nA[Raw Data] --> B[Preprocessing]\\nB --> C[XGBoost]",
-      "mlopsBestPractices": ["Use DVC for data", "Monitor drift"],
-      "trainingCode": "import xgboost as xgb..."
+      "name": "Neuro-Symbolic Latent Programmer",
+      "description": "A unified architecture combining a ViT (Vision), RoBERTa (Text), and Wav2Vec2 (Audio) encoder into a shared embedding space. A central 'Program Synthesizer' LSTM decodes discrete symbolic tokens (Map, Filter, Join) which are executed by differentiable neural modules. Uses Gumbel-Softmax for end-to-end training of discrete operations.",
+      "pros": ["Interpretable reasoning steps", "Generalizes to new combinations", "End-to-end differentiable"],
+      "cons": ["Unstable training dynamics", "High computational cost during search"],
+      "architectureDiagram": "graph TD\\nA[Image/Text/Audio Input] --> B[Modality Encoders]\\nB --> C[Shared Latent Space]\\nC --> D[Program Synthesizer LSTM]\\nD --> E[Symbolic Tokens]\\nE --> F[Neural Module Network]\\nF --> G[Execution Result]\\nG --> H[Loss Calculation]",
+      "mlopsBestPractices": [
+        "Data: WebDataset for sharded multimodal data",
+        "Training: Distributed Data Parallel (DDP) on A100 cluster",
+        "Monitoring: Track discrete token distribution entropy",
+        "Versioning: Git LFS for checkpoints"
+      ],
+      "trainingCode": "import torch\\nimport torch.nn as nn\\n\\nclass NeuroSymbolicModel(nn.Module):\\n    def __init__(self):\\n        super().__init__()\\n        self.vision_enc = VisionTransformer()\\n        self.text_enc = BERT()\\n        # ..."
     }
   ]
 }
 """
 
 def generate_solution(description):
-    prompt = f"{SYSTEM_PROMPT}\n\n{FEW_SHOT_EXAMPLES}\n\nUser Input: \"{description}\"\n\nJSON Response:"
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": f"User Input: \"{description}\"\n\nProvide the JSON solution:"}
+    ]
     
     try:
-        response = client.text_generation(
-            prompt,
-            max_new_tokens=2048,
+        # Use chat_completion for conversational models (Zephyr, Mistral, Llama)
+        response = client.chat_completion(
+            messages,
+            max_tokens=2500,
             temperature=0.7,
-            top_p=0.95,
-            return_full_text=False
+            top_p=0.95
         )
         
-        # clean up response to ensure it's valid JSON
-        json_str = response.strip()
-        if json_str.startswith("```json"):
+        content = response.choices[0].message.content
+        
+        # Robust JSON cleanup
+        json_str = content.strip()
+        if "```json" in json_str:
             json_str = json_str.split("```json")[1].split("```")[0].strip()
-        elif json_str.startswith("```"):
-             json_str = json_str.split("```")[1].split("```")[0].strip()
-             
+        elif "```" in json_str:
+            json_str = json_str.split("```")[1].split("```")[0].strip()
+            
         return json_str
     except Exception as e:
         return json.dumps({
             "error": str(e),
-            "analysis": {"dataType": "text", "taskType": "nlp", "complexity": "low", "domain": "general"},
+            "analysis": {"dataType": "text", "taskType": "nlp", "complexity": "low", "domain": "error_fallback"},
             "recommendations": []
         })
 
 demo = gr.Interface(
     fn=generate_solution,
-    inputs=gr.Textbox(lines=5, placeholder="Describe your ML problem..."),
+    inputs=gr.Textbox(lines=5, placeholder="Describe your complex ML problem..."),
     outputs=gr.JSON(label="Recommendations"),
-    title="ModelForge AI Backend",
-    description="Generates ML recommendations via LLM."
+    title="ModelForge AI Backend v2",
+    description="Advanced ML Solutions Architect (Powered by Mistral-7B-Instruct)"
 )
 
 if __name__ == "__main__":