Update app.py

app.py

@@ -1,138 +1,205 @@
-from transformers import pipeline, AutoTokenizer
 import gradio as gr
 import torch
 import time
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import plotly.express as px
+from transformers import pipeline, AutoTokenizer
+from io import BytesIO
+import base64
+import warnings
+warnings.filterwarnings("ignore")
 
-# ===== RELIABLE MODEL LOADING =====
-def load_model():
-    """Guaranteed model loading with multiple fallbacks"""
-    device = 0 if torch.cuda.is_available() else -1
-    dtype = torch.float16 if device == 0 else torch.float32
-    
-    # Try multiple models in order
-    models = [
-        ("mistralai/Mistral-7B-v0.1", {}),  # Open-access
-        ("google/gemma-2b-it", {"low_cpu_mem_usage": True})  # Gated
-    ]
-    
-    for model_name, kwargs in models:
+# ===== CORE SYSTEM =====
+class AISystem:
+    def __init__(self):
+        self.device = 0 if torch.cuda.is_available() else -1
+        self.dtype = torch.float16 if self.device == 0 else torch.float32
+        self.model = None
+        self.tokenizer = None
+        self.load_models()
+
+    def load_models(self):
+        """Smart model loading with multiple fallbacks"""
+        models = [
+            ("mistralai/Mistral-7B-v0.1", {}),  # Open-access
+            ("google/gemma-2b-it", {"low_cpu_mem_usage": True})  # Gated
+        ]
+        
+        for model_name, kwargs in models:
+            try:
+                self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+                self.model = pipeline(
+                    "text-generation",
+                    model=model_name,
+                    tokenizer=self.tokenizer,
+                    device=self.device,
+                    torch_dtype=self.dtype,
+                    **kwargs
+                )
+                # Verify model works
+                test_output = self.generate("Test", simple=True)
+                if test_output and len(test_output.split()) > 3:
+                    print(f"✅ Loaded {model_name}")
+                    return
+            except Exception as e:
+                print(f"⚠️ Failed {model_name}: {str(e)}")
+        
+        raise RuntimeError("All models failed to load")
+
+    def generate(self, prompt, simple=False):
+        """Guaranteed generation with error handling"""
         try:
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = pipeline(
-                "text-generation",
-                model=model_name,
-                tokenizer=tokenizer,
-                device=device,
-                torch_dtype=dtype,
-                **kwargs
-            )
-            # Test generation
-            test_output = model("Test", max_new_tokens=10)[0]['generated_text']
-            if test_output.strip():
-                print(f"✅ Loaded {model_name}")
-                return model, tokenizer
-        except Exception as e:
-            print(f"⚠️ Failed {model_name}: {str(e)}")
-    
-    raise RuntimeError("All models failed to load")
+            full_prompt = prompt if simple else f"""
+            Provide a detailed, step-by-step answer. Include graphs if requested.
+
+            Question: {prompt}
 
-# Initialize with error handling
+            Answer:
+            1."""
+            
+            output = self.model(
+                full_prompt,
+                max_new_tokens=250,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id
+            )[0]['generated_text']
+            
+            return output.split("Answer:")[-1].strip()
+        except Exception:
+            return "I couldn't generate a response. Please try again."
+
+    def create_graph(self, data_type):
+        """Generate different graph types"""
+        try:
+            x = np.linspace(0, 10, 100)
+            if data_type == "linear":
+                y = x
+                plt.plot(x, y)
+                plt.title("Linear Relationship")
+            elif data_type == "quadratic":
+                y = x**2
+                plt.plot(x, y)
+                plt.title("Quadratic Relationship")
+            elif data_type == "random":
+                y = np.random.rand(100)
+                plt.scatter(x, y)
+                plt.title("Random Data")
+            
+            plt.xlabel("X-axis")
+            plt.ylabel("Y-axis")
+            buf = BytesIO()
+            plt.savefig(buf, format='png')
+            plt.close()
+            return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
+        except Exception:
+            return None
+
+# Initialize system
 try:
-    model, tokenizer = load_model()
+    ai_system = AISystem()
 except Exception as e:
-    model = None
-    print(f"🔴 Critical error: {str(e)}")
+    print(f"🔴 System initialization failed: {str(e)}")
+    ai_system = None
 
-# ===== GUARANTEED GENERATION =====
-def generate_answer(prompt):
-    """Always returns a meaningful answer"""
+# ===== GRADIO INTERFACE =====
+def process_query(prompt):
+    start_time = time.time()
+    
     if not prompt.strip():
-        return "Please enter a valid question"
+        return "Please enter a valid question", None
     
-    if model is None:
-        return "System error - please try again later"
+    if ai_system is None:
+        return "System initialization failed - please check logs", None
     
-    try:
-        start_time = time.time()
-        
-        # Robust prompt engineering
-        full_prompt = f"""Provide a detailed step-by-step answer to this question:
-
-Question: {prompt}
-
-Answer in clear numbered steps:
-1."""
-        
-        output = model(
-            full_prompt,
-            max_new_tokens=200,
-            temperature=0.7,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )[0]['generated_text']
-        
-        # Extract and format answer
-        answer = output.split("Answer in clear numbered steps:")[-1]
-        answer = answer.strip()
-        
-        # Ensure we got actual content
-        if not answer or len(answer.split()) < 3:
-            answer = "I couldn't generate a proper answer. Please try rephrasing your question."
-        
-        # Calculate metrics
-        gen_time = time.time() - start_time
-        word_count = len(answer.split())
-        
-        return f"""📚 Step-by-Step Answer:
+    # Check for graph requests
+    graph_type = None
+    graph_keywords = {
+        "linear graph": "linear",
+        "quadratic graph": "quadratic",
+        "random data": "random",
+        "plot": "linear",
+        "chart": "linear"
+    }
+    
+    for keyword, g_type in graph_keywords.items():
+        if keyword in prompt.lower():
+            graph_type = g_type
+            break
+    
+    # Generate response
+    response = ai_system.generate(prompt)
+    
+    # Create graph if requested
+    graph = None
+    if graph_type:
+        graph = ai_system.create_graph(graph_type)
+    
+    # Format output
+    gen_time = time.time() - start_time
+    formatted_response = f"""📊 Step-by-Step Answer:
 
-{answer}
+{response}
 
-⏱️ Generated in {gen_time:.2f}s | {word_count} words"""
+⏱️ Generated in {gen_time:.2f} seconds"""
     
-    except Exception as e:
-        return f"Error generating answer: {str(e)}"
+    return formatted_response, graph
 
-# ===== COMPLETE UI =====
-with gr.Blocks(theme=gr.themes.Default(), title="🔍 Expert Answer Bot") as demo:
-    gr.Markdown("""<h1><center>Get Detailed Explanations</center></h1>""")
+with gr.Blocks(theme=gr.themes.Soft(), title="🧠 AI Expert Assistant") as demo:
+    gr.Markdown("""<h1><center>Intelligent Answer Engine</center></h1>""")
     
     with gr.Row():
-        question = gr.Textbox(
+        query = gr.Textbox(
             label="Your Question",
-            placeholder="How does blockchain technology work?",
+            placeholder="Ask anything... (e.g. 'Explain photosynthesis and show a linear graph')",
             lines=3
         )
     
     with gr.Row():
-        submit_btn = gr.Button("Get Answer", variant="primary")
+        submit_btn = gr.Button("Generate Answer", variant="primary")
     
     with gr.Row():
         answer = gr.Textbox(
-            label="Step-by-Step Explanation",
+            label="Detailed Explanation",
             lines=8,
             interactive=False
         )
     
-    # Examples that are known to work
+    with gr.Row():
+        graph_output = gr.Image(
+            label="Generated Graph",
+            visible=False
+        )
+    
+    # Example queries
     gr.Examples(
         examples=[
-            "Explain how photosynthesis works in plants",
-            "Describe the steps to solve a quadratic equation",
-            "How does a neural network learn? List the steps"
+            "Explain quantum computing and show a linear graph",
+            "Describe the water cycle with a quadratic graph",
+            "How does machine learning work? Show random data"
         ],
-        inputs=question
+        inputs=query
     )
     
+    def update_ui(response, graph):
+        if graph:
+            return response, gr.update(visible=True, value=graph)
+        return response, gr.update(visible=False)
+    
     submit_btn.click(
-        fn=generate_answer,
-        inputs=question,
-        outputs=answer
+        fn=process_query,
+        inputs=query,
+        outputs=[answer, graph_output]
+    ).then(
+        fn=update_ui,
+        inputs=[answer, graph_output],
+        outputs=[answer, graph_output]
     )
 
-# ===== FAILSAFE LAUNCH =====
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
+        server_port=7860
     )