Update app.py

app.py

@@ -16,6 +16,12 @@ from huggingface_hub import InferenceClient
 import io
 import mimetypes
 import base64
+import cv2
+import numpy as np
+from io import BytesIO
+import tempfile
+import subprocess
+import sys
 
 # -------------------------
 # Environment & constants
@@ -72,17 +78,60 @@ def retry_hf_inference(func):
 
 @retry_hf_inference
 def image_qa_bytes(data: bytes, prompt: str) -> str:
-    """Query MiniGPT-4-V for image-based QA using bytes."""
+    """Query LLaVA for image-based QA using bytes."""
     headers = {"Content-Type": "application/octet-stream"}
-    return client.post("Vision-CAIR/MiniGPT4-V", data=data, headers=headers)
+    return client.post("llava-hf/llava-v1.6-mistral-7b-hf", data=data, headers=headers)
 
 @retry_hf_inference
 def video_label_bytes(data: bytes) -> str:
-    """Get video classification using VideoMAE-Base-Short from bytes."""
+    """Get video classification using VideoMAE-Base from bytes."""
+    # Process video to get first 8 seconds, 16 frames
+    
+    # Read video from bytes
+    video_bytes = BytesIO(data)
+    cap = cv2.VideoCapture()
+    cap.open(video_bytes)
+    
+    # Get video properties
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+    # Calculate frames to extract (16 frames over 8 seconds)
+    target_frames = 16
+    target_duration = 8  # seconds
+    frame_interval = max(1, int(frame_count / (fps * target_duration)))
+    
+    frames = []
+    frame_idx = 0
+    
+    while len(frames) < target_frames and frame_idx < frame_count:
+        ret, frame = cap.read()
+        if not ret:
+            break
+            
+        if frame_idx % frame_interval == 0:
+            # Resize frame to match VideoMAE's expected input
+            frame = cv2.resize(frame, (224, 224))
+            frames.append(frame)
+            
+        frame_idx += 1
+    
+    cap.release()
+    
+    # If we don't have enough frames, duplicate the last frame
+    while len(frames) < target_frames:
+        frames.append(frames[-1])
+    
+    # Stack frames and convert to bytes
+    video_array = np.stack(frames)
+    _, buffer = cv2.imencode('.mp4', video_array)
+    processed_bytes = buffer.tobytes()
+    
+    # Send to VideoMAE
     headers = {"Content-Type": "application/octet-stream"}
     preds = client.post(
-        "MCG-NJU/videomae-base-short-finetuned-ucf101", 
-        data=data,
+        "MCG-NJU/videomae-base-finetuned-ucf101", 
+        data=processed_bytes,
         headers=headers
     )
     return sorted(preds, key=lambda x: x["score"], reverse=True)[0]["label"]
@@ -100,6 +149,18 @@ def sheet_answer_bytes(data: bytes, question: str) -> str:
     label = df.columns[col]
     return f"{label}: {value}"
 
+# -------------------------
+# Code Analysis helpers
+# -------------------------
+
+def run_python(code: str) -> str:
+    """Quick & dirty evaluator for Python code."""
+    with tempfile.NamedTemporaryFile("w+", suffix=".py", delete=False) as f:
+        f.write(code)
+        f.flush()
+        out = subprocess.check_output([sys.executable, f.name], timeout=10)
+    return out.decode().strip()
+
 # -------------------------
 # State definition
 # -------------------------
@@ -241,6 +302,47 @@ class BasicAgent:
         state["current_step"] = "answer"
         return state
 
+    def _code_analysis_node(self, state: AgentState) -> AgentState:
+        """Handle code analysis questions."""
+        try:
+            # Extract code from the question
+            code_match = re.search(r'```python\n(.*?)\n```', state["question"], re.DOTALL)
+            if not code_match:
+                state["logs"]["code_error"] = "No Python code found in question"
+                state["current_step"] = "answer"
+                return state
+            
+            code = code_match.group(1)
+            
+            # Run the code and get output
+            try:
+                output = run_python(code)
+                state["history"].append({
+                    "step": "code",
+                    "output": f"Code execution result:\n{output}"
+                })
+            except subprocess.TimeoutExpired:
+                state["history"].append({
+                    "step": "code",
+                    "output": "Code execution timed out after 10 seconds"
+                })
+            except subprocess.CalledProcessError as e:
+                state["history"].append({
+                    "step": "code",
+                    "output": f"Code execution failed with error:\n{e.output.decode()}"
+                })
+            except Exception as e:
+                state["history"].append({
+                    "step": "code",
+                    "output": f"Unexpected error during code execution:\n{str(e)}"
+                })
+            
+        except Exception as e:
+            state["logs"]["code_error"] = str(e)
+        
+        state["current_step"] = "answer"
+        return state
+
     def _generate_answer(self, state: AgentState) -> AgentState:
         # Collect all tool outputs with clear section headers
         materials = []
@@ -349,6 +451,7 @@ Write ANSWER: <answer> on its own line.
         sg.add_node("image", self._image_node)
         sg.add_node("video", self._video_node)
         sg.add_node("sheet", self._sheet_node)
+        sg.add_node("code", self._code_analysis_node)
 
         # Add edges
         sg.add_edge("analyze", "search")
@@ -357,6 +460,7 @@ Write ANSWER: <answer> on its own line.
         sg.add_edge("image", "answer")
         sg.add_edge("video", "answer")
         sg.add_edge("sheet", "answer")
+        sg.add_edge("code", "answer")
 
         def router(state: AgentState):
             return state["current_step"]
@@ -366,7 +470,8 @@ Write ANSWER: <answer> on its own line.
             "answer": "answer",
             "image": "image",
             "video": "video",
-            "sheet": "sheet"
+            "sheet": "sheet",
+            "code": "code"
         })
         
         sg.set_entry_point("analyze")
@@ -465,36 +570,22 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     
     for item in questions_data:
         task_id = item.get("task_id")
-        question_text = item.get("question")
-        if not task_id or question_text is None:
-            print(f"Skipping item with missing task_id or question: {item}")
+        if not task_id:
+            print(f"Skipping item with missing task_id: {item}")
             continue
 
         try:
-            print(f"\nProcessing question {task_id}: {question_text[:50]}...")
-            
-            # Initialize state for this question
-            state: AgentState = {
-                "question": question_text,
-                "current_step": "analyze",
-                "final_answer": "",
-                "history": [],
-                "needs_search": False,
-                "search_query": "",
-                "task_id": task_id,
-                "logs": {},
-                "file_url": ""
-            }
+            print(f"\nProcessing question {task_id}...")
             
-            # Run the workflow
-            final_state = agent.workflow.invoke(state)
-            answer = final_state["final_answer"]
+            # Pass the entire item as JSON string
+            state_question = json.dumps(item)
+            answer = agent(state_question, task_id)
             
             # Add to results
             answers_payload.append({"task_id": task_id, "submitted_answer": answer})
             results_log.append({
                 "Task ID": task_id,
-                "Question": question_text,
+                "Question": item.get("question", ""),
                 "Submitted Answer": answer
             })
             
@@ -504,7 +595,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
             print(f"Error running agent on task {task_id}: {e}")
             results_log.append({
                 "Task ID": task_id,
-                "Question": question_text,
+                "Question": item.get("question", ""),
                 "Submitted Answer": f"ERROR: {e}"
             })