Update app.py

app.py

@@ -22,92 +22,29 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # --- Basic Agent Definition ---
 
 
-
 class BasicAgent:
-    def __init__(self, model="google/gemma-7b"):
+    def __init__(self, model="google/gemma-2b"):  # Smaller 2B version recommended
         self.tokenizer = AutoTokenizer.from_pretrained(model)
-        self.model = AutoModelForCausalLM.from_pretrained(model)
-        print("BasicAgent initialized with AutoModel")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model,
+            device_map="auto",
+            torch_dtype=torch.float32,  # Explicitly use float32 for CPU
+            low_cpu_mem_usage=True      # Reduces memory spikes
+        )
+        print(f"Initialized on device: {self.model.device}")
 
-    def __call__(self, question: str) -> str:
-        inputs = self.tokenizer(question, return_tensors="pt")
-        outputs = self.model.generate(**inputs, max_new_tokens=100)
-        return self.tokenizer.decode(outputs[0])
+    def __call__(self, question: str, max_tokens: int = 100) -> str:
+        inputs = self.tokenizer(question, return_tensors="pt").to(self.model.device)
+        with torch.no_grad():  # Reduces memory usage
+            outputs = self.model.generate(**inputs, max_new_tokens=max_tokens)
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             
     def wikipedia_search(self, query: str) -> str:
         """Get Wikipedia summary"""
         page = self.wiki.page(query)
         return page.summary if page.exists() else "No Wikipedia page found"
 
-    def process_document(self, file_path: str) -> str:
-        """Handle PDF, Word, CSV, Excel files"""
-        if not os.path.exists(file_path):
-            return "File not found"
-        
-        ext = os.path.splitext(file_path)[1].lower()
-        
-        try:
-            if ext == '.pdf':
-                return self._process_pdf(file_path)
-            elif ext in ('.doc', '.docx'):
-                return self._process_word(file_path)
-            elif ext == '.csv':
-                return pd.read_csv(file_path).to_string()
-            elif ext in ('.xls', '.xlsx'):
-                return pd.read_excel(file_path).to_string()
-            else:
-                return "Unsupported file format"
-        except Exception as e:
-            return f"Error processing document: {str(e)}"
-
-    def _process_pdf(self, file_path: str) -> str:
-        """Process PDF using Gemini's vision capability"""
-        try:
-            # For Gemini 1.5 or later which supports file uploads
-            with open(file_path, "rb") as f:
-                file = genai.upload_file(f)
-                response = self.model.generate_content(
-                    ["Extract and summarize the key points from this document:", file]
-                )
-                return response.text
-        except:
-            # Fallback for older Gemini versions
-            try:
-                import PyPDF2
-                with open(file_path, 'rb') as f:
-                    reader = PyPDF2.PdfReader(f)
-                    return "\n".join([page.extract_text() for page in reader.pages])
-            except ImportError:
-                return "PDF processing requires PyPDF2 (pip install PyPDF2)"
-
-    def _process_word(self, file_path: str) -> str:
-        """Process Word documents"""
-        try:
-            from docx import Document
-            doc = Document(file_path)
-            return "\n".join([para.text for para in doc.paragraphs])
-        except ImportError:
-            return "Word processing requires python-docx (pip install python-docx)"
 
-    def process_request(self, request: Union[str, Dict]) -> str:
-        """
-        Handle different request types:
-        - Direct text queries
-        - File processing requests
-        - Complex multi-step requests
-        """
-        if isinstance(request, dict):
-            if 'steps' in request:
-                results = []
-                for step in request['steps']:
-                    if step['type'] == 'search':
-                        results.append(self.web_search(step['query']))
-                    elif step['type'] == 'process':
-                        results.append(self.process_document(step['file']))
-                return self.generate_response(f"Process these results: {results}")
-            return "Unsupported request format"
-        
-        return self.generate_response(request)
 
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):