initial impl

app.py

@@ -3,6 +3,13 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
+from llama_index import VectorStoreIndex, SimpleDirectoryReader, Document
+from llama_index.llms import HuggingFaceLLM
+from llama_index import ServiceContext
+from llama_index.embeddings import HuggingFaceEmbedding
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
+from dotenv import load_dotenv
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -13,11 +20,94 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
+        load_dotenv()
+        
+        # Initialize model and tokenizer 
+        model_name = "open-r1/OlympicCoder-7B"
+        
+        # Configure quantization
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+        )
+        
+        # Initialize tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            quantization_config=quantization_config,
+            trust_remote_code=True
+        )
+        
+        # Create HuggingFaceLLM instance
+        self.llm = HuggingFaceLLM(
+            context_window=4096,
+            max_new_tokens=512,
+            tokenizer=self.tokenizer,
+            model=self.model,
+            model_name=model_name,
+            device_map="auto",
+            temperature=0.1
+        )
+        
+        print("Agent initialized.")
+
     def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        fixed_answer = "This is a default answer."
-        print(f"Agent returning fixed answer: {fixed_answer}")
-        return fixed_answer
+        print(f"Agent received question: {question[:50]}...")
+        
+        try:
+            # Extract data from the question if present
+            data_start = question.find('```python')
+            if data_start != -1:
+                data_end = question.find('```', data_start + 8)
+                data_code = question[data_start + 8:data_end].strip()
+                
+                # Execute the data code in a safe context to create DataFrame
+                local_vars = {}
+                exec(data_code, {"pd": pd}, local_vars)
+                df = local_vars.get('df')
+                
+                # Extract the actual question (usually after the code block)
+                actual_question = question[data_end + 3:].strip()
+                
+                # Create analysis prompt
+                prompt = f"""
+                You are a data analysis expert. Given the following DataFrame and question, 
+                provide the correct answer using pandas operations.
+                
+                DataFrame head:
+                {df.head().to_string()}
+                
+                DataFrame info:
+                {df.info()}
+                
+                Question: {actual_question}
+                
+                Provide only the final answer without explanations.
+                """
+                
+            else:
+                # Handle non-data questions
+                prompt = f"""
+                You are a helpful AI assistant specializing in data analysis.
+                Please provide a clear, concise answer to this question:
+                {question}
+                """
+            
+            # Use LlamaIndex query engine
+            query_engine = self.index.as_query_engine()
+            response = query_engine.query(prompt)
+            answer = str(response)
+            
+            print(f"Agent generated answer: {answer[:100]}...")
+            return answer
+            
+        except Exception as e:
+            print(f"Error generating answer: {e}")
+            return f"I apologize, I encountered an error: {str(e)}"
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """

requirements.txt

@@ -1,2 +1,10 @@
 gradio
-requests
+requests
+llama-index
+transformers
+python-dotenv
+torch
+sentence-transformers
+nltk
+accelerate
+bitsandbytes