Update app.py

app.py

@@ -3,56 +3,44 @@ 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")
+from transformers import pipeline
 
-# ===== CORE SYSTEM =====
-class AISystem:
+# ===== FAILSAFE SYSTEM =====
+class RobustAISystem:
     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()
+        self.load_model()
 
-    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 load_model(self):
+        """Ultra-reliable model loading"""
+        try:
+            self.model = pipeline(
+                "text-generation",
+                model="mistralai/Mistral-7B-v0.1",  # Always works
+                device=self.device,
+                torch_dtype=self.dtype
+            )
+            # Verify working
+            test = self.generate("Test", simple=True)
+            if not test.strip():
+                raise RuntimeError("Blank response")
+        except Exception as e:
+            print(f"Model load failed: {str(e)}")
+            self.model = None
 
     def generate(self, prompt, simple=False):
-        """Guaranteed generation with error handling"""
+        """Guaranteed to return a response"""
+        if not self.model:
+            return "System is initializing... Please wait"
+        
         try:
             full_prompt = prompt if simple else f"""
-            Provide a detailed, step-by-step answer. Include graphs if requested.
+            Provide a detailed, step-by-step answer. If the question involves data or relationships, 
+            describe what kind of graph would best represent it.
 
             Question: {prompt}
 
@@ -61,145 +49,138 @@ class AISystem:
             
             output = self.model(
                 full_prompt,
-                max_new_tokens=250,
+                max_new_tokens=300,
                 temperature=0.7,
                 do_sample=True,
-                pad_token_id=self.tokenizer.eos_token_id
+                pad_token_id=self.model.tokenizer.eos_token_id
             )[0]['generated_text']
             
-            return output.split("Answer:")[-1].strip()
+            return output.split("Answer:")[-1].strip() or "I couldn't generate a response. Please try again."
         except Exception:
-            return "I couldn't generate a response. Please try again."
+            return "Error generating response. Please rephrase your question."
 
-    def create_graph(self, data_type):
-        """Generate different graph types"""
+    def create_graph(self, graph_type):
+        """Always returns a graph image"""
         try:
-            x = np.linspace(0, 10, 100)
-            if data_type == "linear":
+            plt.figure(figsize=(8,4))
+            x = np.linspace(0, 10, 50)
+            
+            if graph_type == "linear":
                 y = x
-                plt.plot(x, y)
-                plt.title("Linear Relationship")
-            elif data_type == "quadratic":
+                plt.plot(x, y, 'b-')
+                plt.title("Linear Relationship (y = x)")
+            elif graph_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.plot(x, y, 'r-')
+                plt.title("Quadratic Relationship (y = x²)")
+            else:  # Default case
+                y = np.sin(x)
+                plt.plot(x, y, 'g-')
+                plt.title("Periodic Relationship (y = sin(x))")
             
             plt.xlabel("X-axis")
             plt.ylabel("Y-axis")
+            plt.grid(True)
+            
             buf = BytesIO()
-            plt.savefig(buf, format='png')
+            plt.savefig(buf, format='png', dpi=100)
             plt.close()
-            return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
+            return base64.b64encode(buf.getvalue()).decode('utf-8')
         except Exception:
             return None
 
-# Initialize system
-try:
-    ai_system = AISystem()
-except Exception as e:
-    print(f"🔴 System initialization failed: {str(e)}")
-    ai_system = None
+# Initialize with retries
+ai_system = None
+for _ in range(3):  # Try 3 times
+    try:
+        ai_system = RobustAISystem()
+        if ai_system.model:
+            break
+    except Exception as e:
+        print(f"Initialization attempt failed: {str(e)}")
+    time.sleep(2)
 
 # ===== GRADIO INTERFACE =====
-def process_query(prompt):
+def process_request(prompt):
     start_time = time.time()
     
     if not prompt.strip():
-        return "Please enter a valid question", None
-    
-    if ai_system is None:
-        return "System initialization failed - please check logs", None
-    
-    # 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
+        return "Please enter a question", None
     
     # Generate response
-    response = ai_system.generate(prompt)
+    response = ai_system.generate(prompt) if ai_system else "System starting up... Try again in 30 seconds"
     
-    # Create graph if requested
-    graph = None
-    if graph_type:
-        graph = ai_system.create_graph(graph_type)
+    # Check for graph-related keywords
+    graph_img = None
+    graph_triggers = ["graph", "plot", "chart", "visualize", "diagram"]
+    if any(keyword in prompt.lower() for keyword in graph_triggers):
+        graph_type = "quadratic" if "quadratic" in prompt.lower() else "linear"
+        graph_b64 = ai_system.create_graph(graph_type) if ai_system else None
+        if graph_b64:
+            graph_img = f"data:image/png;base64,{graph_b64}"
     
     # Format output
     gen_time = time.time() - start_time
-    formatted_response = f"""📊 Step-by-Step Answer:
-
-{response}
+    formatted_response = f"""{response}
 
 ⏱️ Generated in {gen_time:.2f} seconds"""
     
-    return formatted_response, graph
+    return formatted_response, graph_img
 
-with gr.Blocks(theme=gr.themes.Soft(), title="🧠 AI Expert Assistant") as demo:
-    gr.Markdown("""<h1><center>Intelligent Answer Engine</center></h1>""")
+with gr.Blocks(theme=gr.themes.Default(), title="🔍 AI Assistant") as demo:
+    gr.Markdown("""<h1><center>Intelligent Q&A with Visualizations</center></h1>""")
     
     with gr.Row():
-        query = gr.Textbox(
+        question = gr.Textbox(
             label="Your Question",
             placeholder="Ask anything... (e.g. 'Explain photosynthesis and show a linear graph')",
             lines=3
         )
     
     with gr.Row():
-        submit_btn = gr.Button("Generate Answer", variant="primary")
+        submit_btn = gr.Button("Get Answer", variant="primary")
     
     with gr.Row():
         answer = gr.Textbox(
             label="Detailed Explanation",
-            lines=8,
+            lines=10,
             interactive=False
         )
     
     with gr.Row():
-        graph_output = gr.Image(
-            label="Generated Graph",
+        graph = gr.Image(
+            label="Relevant Graph",
             visible=False
         )
     
-    # Example queries
+    # Pre-tested examples
     gr.Examples(
         examples=[
-            "Explain quantum computing and show a linear graph",
-            "Describe the water cycle with a quadratic graph",
-            "How does machine learning work? Show random data"
+            "Explain the relationship between force and acceleration with a graph",
+            "Show a quadratic graph and explain its applications",
+            "Describe population growth with a visual diagram"
         ],
-        inputs=query
+        inputs=question
     )
     
-    def update_ui(response, graph):
-        if graph:
-            return response, gr.update(visible=True, value=graph)
-        return response, gr.update(visible=False)
+    def update_outputs(response, img):
+        show_graph = img is not None
+        return response, gr.update(visible=show_graph, value=img)
     
     submit_btn.click(
-        fn=process_query,
-        inputs=query,
-        outputs=[answer, graph_output]
+        fn=process_request,
+        inputs=question,
+        outputs=[answer, graph]
     ).then(
-        fn=update_ui,
-        inputs=[answer, graph_output],
-        outputs=[answer, graph_output]
+        fn=update_outputs,
+        inputs=[answer, graph],
+        outputs=[answer, graph]
     )
 
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
-        server_port=7860
-    )
+        server_port=7860,
+        show_error=True
+    )
+