Upload 6 files

backend/extractors.py

@@ -0,0 +1,57 @@
+from pathlib import Path
+import fitz  # PyMuPDF
+from docx import Document
+from pptx import Presentation
+
+# PDF
+def extract_text_pdf(file_path: str) -> str:
+    text = ""
+    with fitz.open(file_path) as doc:
+        for page in doc:
+            text += page.get_text()
+    return text.strip()
+
+# DOCX
+def extract_text_docx(file_path: str) -> dict:
+    doc = Document(file_path)
+    paragraphs = []
+
+    for para in doc.paragraphs:
+        text = para.text.strip()
+        if text:  # Only include non-empty paragraphs
+            paragraphs.append({
+                "style": para.style.name,
+                "text": text
+            })
+
+    return {"content": paragraphs}
+
+# TXT
+def extract_text_txt(file_path: str) -> str:
+    with open(file_path, "r", encoding="utf-8") as f:
+        lines = f.read().splitlines()  # ✅ split into clean lines
+    return {"content": lines}
+
+# PPTX
+def extract_text_pptx(file_path: str) -> str:
+    prs = Presentation(file_path)
+    text = []
+    for slide in prs.slides:
+        for shape in slide.shapes:
+            if hasattr(shape, "text"):
+                text.append(shape.text)
+    return "\n".join(text).strip()
+
+# Dispatcher
+def extract_text(file_path: str) -> str:
+    ext = Path(file_path).suffix.lower()
+    if ext == ".pdf":
+        return extract_text_pdf(file_path)
+    elif ext == ".docx":
+        return extract_text_docx(file_path)
+    elif ext == ".txt":
+        return extract_text_txt(file_path)
+    elif ext == ".pptx":
+        return extract_text_pptx(file_path)
+    else:
+        raise ValueError("Unsupported file extension")

backend/file_handler.py

@@ -0,0 +1,49 @@
+from fastapi import UploadFile, HTTPException
+from pathlib import Path
+from uuid import uuid4
+
+# Accepted file types
+ALLOWED_TYPES = {
+    # Document types
+    "application/pdf": ".pdf",
+    "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
+    "text/plain": ".txt",
+    "application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx",
+    
+    # Image types
+    "image/jpeg": ".jpg",
+    "image/jpg": ".jpg",
+    "image/png": ".png",
+    "image/gif": ".gif",
+    "image/bmp": ".bmp",
+    "image/webp": ".webp"
+}
+
+MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
+
+tmp_dir = Path("/tmp/uploads")
+tmp_dir.mkdir(parents=True, exist_ok=True)
+
+UPLOAD_DIR = tmp_dir
+
+def save_upload(file: UploadFile) -> tuple[str, str]:
+    # Print file content type for debugging
+    print(f"File content type: {file.content_type}")
+    
+    if file.content_type not in ALLOWED_TYPES:
+        raise HTTPException(status_code=400, detail=f"Unsupported file type: {file.content_type}")
+
+    # Read file into memory to check size and save it
+    file_bytes = file.file.read()
+
+    if len(file_bytes) > MAX_FILE_SIZE:
+        raise HTTPException(status_code=413, detail="File is too large. Maximum size is 10MB.")
+
+    file_ext = ALLOWED_TYPES[file.content_type]
+    file_id = str(uuid4())
+    file_path = UPLOAD_DIR / f"{file_id}{file_ext}"
+
+    with open(file_path, "wb") as f:
+        f.write(file_bytes)
+
+    return file_id, file_path.name

backend/image_processor.py

@@ -0,0 +1,101 @@
+from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+from transformers import ViltProcessor, ViltForQuestionAnswering  # Add these imports
+from PIL import Image
+import torch
+from typing import Dict, Any, Union, List
+
+class ImageProcessor:
+    def __init__(self, caption_model_name: str = "nlpconnect/vit-gpt2-image-captioning", 
+                 vqa_model_name: str = "dandelin/vilt-b32-finetuned-vqa"):
+        # Image captioning model
+        self.caption_processor = ViTImageProcessor.from_pretrained(caption_model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(caption_model_name)
+        self.caption_model = VisionEncoderDecoderModel.from_pretrained(caption_model_name)
+        
+        # VQA model
+        self.vqa_processor = ViltProcessor.from_pretrained(vqa_model_name)
+        self.vqa_model = ViltForQuestionAnswering.from_pretrained(vqa_model_name)
+        
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.caption_model.to(self.device)
+        self.vqa_model.to(self.device)
+        
+    def generate_caption(self, image_path: str) -> str:
+        """Generate a descriptive caption for the provided image"""
+        try:
+            image = Image.open(image_path).convert("RGB")
+            pixel_values = self.caption_processor(image, return_tensors="pt").pixel_values.to(self.device)
+            
+            gen_kwargs = {
+                "max_length": 50,
+                "num_beams": 4,
+                "early_stopping": True
+            }
+            
+            output_ids = self.caption_model.generate(pixel_values, **gen_kwargs)
+            caption = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+            
+            return caption
+        except Exception as e:
+            return f"Error processing image: {str(e)}"
+    
+    def answer_image_question(self, image_path: str, question: str) -> Dict[str, Any]:
+        """Answer a question about the provided image using a Visual QA model"""
+        try:
+            # Open image
+            image = Image.open(image_path).convert("RGB")
+            
+            # Prepare inputs
+            inputs = self.vqa_processor(image, question, return_tensors="pt")
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            
+            # Forward pass
+            with torch.no_grad():
+                outputs = self.vqa_model(**inputs)
+            
+            # Get answer
+            logits = outputs.logits
+            idx = logits.argmax(-1).item()
+            answer = self.vqa_model.config.id2label[idx]
+            confidence = torch.softmax(logits, dim=-1)[0, idx].item()
+            
+            return {"answer": answer, "confidence": confidence}
+            
+        except Exception as e:
+            print(f"VQA Error: {str(e)}")
+            # Fallback to caption
+            try:
+                caption = self.generate_caption(image_path)
+                return {
+                    "answer": f"Based on the image which shows {caption}, I cannot provide a specific answer.",
+                    "confidence": 0.0
+                }
+            except Exception as e2:
+                return {"answer": f"Error processing image: {str(e)}, {str(e2)}", "confidence": 0.0}
+
+
+# Add this to the end of your image_processor.py file
+
+class SimpleImageProcessor:
+    """A simple fallback image processor that doesn't require external models"""
+    
+    def __init__(self):
+        """Initialize without any models"""
+        print("Using SimpleImageProcessor fallback")
+    
+    def generate_caption(self, image_path: str) -> str:
+        """Generate a basic caption for the provided image"""
+        try:
+            # Just extract basic image information
+            from PIL import Image
+            img = Image.open(image_path)
+            return f"an image of size {img.width}x{img.height}"
+        except Exception as e:
+            return f"an image (could not process: {str(e)})"
+    
+    def answer_image_question(self, image_path: str, question: str) -> Dict[str, Any]:
+        """Provide a fallback answer for image questions"""
+        return {
+            "answer": "I cannot analyze this image right now. The image processing system is not fully functional.",
+            "confidence": 0.0
+        }

backend/models.py

@@ -0,0 +1,129 @@
+from transformers import pipeline, BlipProcessor, Blip2ForConditionalGeneration
+import torch
+import os
+from typing import Dict, Any, Optional
+from PIL import Image
+
+# Singleton model manager
+class ModelManager:
+    _instance = None
+    
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+    
+    def __init__(self):
+        self.pipelines = {}
+        self.specialized_models = {}
+        
+        # Set default models - preferably small ones that are publicly accessible
+        self.model_configs = {
+            "image_captioning": {
+                "name": "Salesforce/blip-image-captioning-base",  # Public model
+                "type": "image-to-text"
+            },
+            "document_qa": {
+                "name": "impira/layoutlm-document-qa",  # Document-specialized model
+                "type": "document-question-answering"
+            },
+            "blip2_vqa": {
+                "name": "Salesforce/blip2-flan-t5-xl",  # BLIP-2 model for improved VQA
+                "type": "specialized"
+            }
+        }
+        
+        # Check if CUDA is available
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    def load_pipeline(self, pipeline_type: str) -> bool:
+        """Load a specific pipeline if it's not already loaded"""
+        if pipeline_type not in self.model_configs:
+            return False
+            
+        if pipeline_type in self.pipelines or pipeline_type in self.specialized_models:
+            return True
+            
+        config = self.model_configs[pipeline_type]
+        model_name = config["name"]
+        
+        try:
+            if config["type"] == "document-question-answering":
+                self.pipelines[pipeline_type] = pipeline("document-question-answering", model=model_name)
+            elif config["type"] == "image-to-text":
+                self.pipelines[pipeline_type] = pipeline("image-to-text", model=model_name)
+            elif config["type"] == "specialized" and pipeline_type == "blip2_vqa":
+                # Load BLIP-2 model components
+                processor = BlipProcessor.from_pretrained(model_name)
+                model = Blip2ForConditionalGeneration.from_pretrained(
+                    model_name, 
+                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
+                )
+                model.to(self.device)
+                
+                self.specialized_models[pipeline_type] = {
+                    "processor": processor,
+                    "model": model
+                }
+                
+            return True
+        except Exception as e:
+            print(f"Error loading pipeline {model_name}: {str(e)}")
+            return False
+    
+    def get_pipeline(self, pipeline_type: str) -> Optional[Any]:
+        """Get a loaded pipeline or load it if not already loaded"""
+        if pipeline_type not in self.pipelines and \
+           pipeline_type not in self.specialized_models and \
+           not self.load_pipeline(pipeline_type):
+            return None
+            
+        if pipeline_type in self.pipelines:
+            return self.pipelines[pipeline_type]
+        else:
+            return self.specialized_models[pipeline_type]
+    
+    def blip2_vqa(self, image_path: str, question: str, max_length: int = 100) -> str:
+        """
+        Perform VQA using BLIP-2 model
+        
+        Args:
+            image_path: Path to the image file
+            question: The question to ask about the image
+            max_length: Maximum length of generated answer
+        
+        Returns:
+            Answer to the question
+        """
+        # Make sure the model is loaded
+        if not self.load_pipeline("blip2_vqa"):
+            return "Model could not be loaded"
+        
+        # Get the model components
+        components = self.specialized_models["blip2_vqa"]
+        processor = components["processor"]
+        model = components["model"]
+        
+        try:
+            # Load and process the image
+            image = Image.open(image_path).convert('RGB')
+            
+            # Process inputs
+            inputs = processor(image, question, return_tensors="pt").to(self.device)
+            
+            # Generate answer
+            with torch.no_grad():
+                generated_ids = model.generate(
+                    **inputs,
+                    max_length=max_length,
+                    do_sample=False,
+                    num_beams=5,
+                )
+            
+            # Decode the generated answer
+            answer = processor.decode(generated_ids[0], skip_special_tokens=True)
+            return answer
+            
+        except Exception as e:
+            return f"Error processing image: {str(e)}"

backend/qa_engine.py

@@ -0,0 +1,65 @@
+from transformers import pipeline
+import torch
+from typing import Dict, List, Any
+
+class QAEngine:
+    def __init__(self, model_name: str = "deepset/roberta-base-squad2"):
+        # Use the pipeline API which works better with Hugging Face Spaces
+        self.qa_pipeline = pipeline("question-answering", model=model_name)
+        
+    def answer_question(self, context: str, question: str) -> Dict[str, Any]:
+        """Answer a question based on the provided context"""
+        try:
+            # Use the pipeline directly
+            result = self.qa_pipeline(question=question, context=context)
+            
+            return {
+                "answer": result["answer"],
+                "confidence": result["score"],
+                "start_position": result["start"],
+                "end_position": result["end"]
+            }
+        except Exception as e:
+            return {
+                "answer": f"Error processing question: {str(e)}",
+                "confidence": 0.0,
+                "start_position": 0,
+                "end_position": 0
+            }
+    
+    def answer_multiple_questions(self, context: str, questions: List[str]) -> List[Dict[str, Any]]:
+        """Answer multiple questions from the same context"""
+        return [self.answer_question(context, question) for question in questions]
+
+
+# Add this to qa_engine.py
+class SimpleQAEngine:
+    """A simple QA engine that doesn't rely on complex models"""
+    
+    def answer_question(self, context: str, question: str) -> Dict[str, Any]:
+        """Basic keyword-based answer extraction (fallback when models fail)"""
+        # Very basic implementation - just finds sentences with keywords from the question
+        from nltk.tokenize import sent_tokenize
+        try:
+            import nltk
+            nltk.download('punkt', quiet=True)
+        except:
+            pass
+            
+        question_words = set(question.lower().split())
+        best_sentence = ""
+        best_score = 0
+        
+        for sentence in sent_tokenize(context):
+            sentence_words = set(sentence.lower().split())
+            overlap = len(question_words.intersection(sentence_words))
+            if overlap > best_score:
+                best_score = overlap
+                best_sentence = sentence
+        
+        return {
+            "answer": best_sentence if best_score > 0 else "No relevant information found.",
+            "confidence": min(best_score / max(1, len(question_words)), 1.0),
+            "start_position": context.find(best_sentence) if best_sentence in context else 0,
+            "end_position": context.find(best_sentence) + len(best_sentence) if best_sentence in context else 0
+        }

backend/response_formatter.py

@@ -0,0 +1,48 @@
+from typing import Dict, Any, List, Union
+
+class ResponseFormatter:
+    @staticmethod
+    def format_document_qa_response(qa_result: Dict[str, Any], document_name: str) -> Dict[str, Any]:
+        """Format the response from the QA engine for document questions"""
+        formatted_response = {
+            "document": document_name,
+            "answer": qa_result.get("answer", "No answer found"),
+            "confidence": round(qa_result.get("confidence", 0) * 100, 2),
+            "metadata": {
+                "source_type": "document",
+                "timestamp": qa_result.get("timestamp")
+            }
+        }
+        
+        # Add highlighted text positions if available
+        if "start_position" in qa_result and "end_position" in qa_result:
+            formatted_response["highlight"] = {
+                "start": qa_result["start_position"],
+                "end": qa_result["end_position"]
+            }
+            
+        return formatted_response
+    
+    @staticmethod
+    def format_image_qa_response(vqa_result: Dict[str, Any], image_name: str) -> Dict[str, Any]:
+        """Format the response from the image QA engine"""
+        formatted_response = {
+            "image": image_name,
+            "answer": vqa_result.get("answer", "No answer found"),
+            "confidence": round(vqa_result.get("confidence", 0) * 100, 2),
+            "metadata": {
+                "source_type": "image",
+                "timestamp": vqa_result.get("timestamp")
+            }
+        }
+        
+        return formatted_response
+    
+    @staticmethod
+    def format_error_response(error_message: str, status_code: int = 400) -> Dict[str, Any]:
+        """Format error responses"""
+        return {
+            "error": True,
+            "message": error_message,
+            "status_code": status_code
+        }