update. focus on remote and pdf processing

app.py

@@ -2,92 +2,58 @@ from flask import Flask, request, jsonify, send_from_directory
 import requests
 import os
 from dotenv import load_dotenv
-import traceback
-from PIL import Image
-from pdf2image import convert_from_bytes
-import base64
-from io import BytesIO
+import logging
 
 load_dotenv()
 
 app = Flask(__name__)
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 API_URL = "https://api-inference.huggingface.co/models/"
 headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_KEY')}"}
-HF_REPO_ID = os.getenv('HF_REPO_ID')  # Your Hugging Face repo ID
+HF_REPO_ID = os.getenv('HF_REPO_ID')
+PDF_FILENAME = os.getenv('PDF_FILENAME')
 
 def query(payload, model):
-    response = requests.post(API_URL + model, headers=headers, json=payload)
-    return response.json()
+    try:
+        response = requests.post(API_URL + model, headers=headers, json=payload)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Error querying model {model}: {str(e)}")
+        return {"error": str(e)}
 
-# Process PDFs using Hugging Face's PDF processing model
-def process_pdfs():
-    pdf_url = f"https://huggingface.co/spaces/{HF_REPO_ID}/resolve/main/data/your_pdf_file.pdf"
+def process_pdf():
+    pdf_url = f"https://huggingface.co/spaces/{HF_REPO_ID}/resolve/main/data/{PDF_FILENAME}"
+    logger.info(f"Attempting to process PDF at URL: {pdf_url}")
     
-    try:
-        # Download PDF
-        pdf_response = requests.get(pdf_url)
-        pdf_response.raise_for_status()  # This will raise an exception for HTTP errors
-        
-        print(f"PDF downloaded successfully. Content length: {len(pdf_response.content)} bytes")
-        
-        # Convert PDF to images
-        images = convert_from_bytes(pdf_response.content)
-        print(f"Converted PDF to {len(images)} images")
-        
-        # Process each image
-        vision_model = "google/vit-base-patch16-224"
-        summaries = []
-        
-        for i, image in enumerate(images):
-            # Convert image to base64
-            buffered = BytesIO()
-            image.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
-            
-            # Process image with vision model
-            payload = {
-                "inputs": [
-                    {
-                        "image": img_str,
-                        "text": "Describe the content of this image in detail."
-                    }
-                ]
-            }
-            response = query(payload, vision_model)
-            print(f"Page {i+1} processing response:", json.dumps(response, indent=2))
-            
-            if isinstance(response, list) and len(response) > 0 and 'generated_text' in response[0]:
-                summaries.append(response[0]['generated_text'])
-            else:
-                summaries.append(f"Error processing page {i+1}")
-        
-        return " ".join(summaries)
-    except Exception as e:
-        print(f"Error in process_pdfs: {str(e)}")
-        print(traceback.format_exc())
-        return f"Error processing PDF: {str(e)}"
+    model = "impira/layoutlm-document-qa"
+    payload = {
+        "inputs": {
+            "question": "Summarize the main points of this document.",
+            "image": pdf_url
+        }
+    }
+    
+    response = query(payload, model)
+    logger.info(f"PDF processing response: {response}")
+    
+    if 'error' in response:
+        return f"Error processing PDF: {response['error']}"
+    elif 'answer' in response:
+        return response['answer']
+    else:
+        return str(response)
 
-# Get the summary of PDFs
-pdf_summary = process_pdfs()
-print("PDF Summary:", pdf_summary)
+# Process PDF and get summary
+pdf_summary = process_pdf()
+logger.info(f"PDF Summary: {pdf_summary}")
 
 # Get embeddings for the summary
 embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
-if not pdf_summary.startswith("Error"):
-    try:
-        summary_embedding = query({"inputs": pdf_summary}, embedding_model)[0]
-        print("Successfully created summary embedding")
-    except Exception as e:
-        print(f"Error getting embedding: {str(e)}")
-        print(traceback.format_exc())
-        summary_embedding = None
-else:
-    print("Skipping embedding due to PDF processing error")
-    summary_embedding = None
-
-if summary_embedding is None:
-    print("WARNING: summary_embedding is None. The chatbot will not be able to provide meaningful responses.")
+summary_embedding = query({"inputs": pdf_summary}, embedding_model)[0] if not pdf_summary.startswith("Error") else None
 
 @app.route('/')
 def home():
@@ -97,17 +63,18 @@ def home():
 def ask():
     prompt = request.json['question']
     
+    if summary_embedding is None:
+        return jsonify({'response': "I'm sorry, but I couldn't process the PDF correctly. Please check the PDF file and try again later."})
+    
     # Get embedding for the question
     query_embedding = query({"inputs": prompt}, embedding_model)[0]
     
     # Calculate similarity
     similarity = sum(a*b for a, b in zip(query_embedding, summary_embedding))
     
-    # Generate character response
-    character_traits = "I am a knowledgeable and friendly AI assistant based on the content of the provided PDFs."
-    input_text = f"Character: {character_traits}\nContext: {pdf_summary}\nUser: {prompt}\nCharacter's response:"
-    
+    # Generate response using T5 model
     generator_model = "google/flan-t5-base"
+    input_text = f"Context: {pdf_summary}\n\nQuestion: {prompt}\n\nAnswer:"
     response = query({"inputs": input_text}, generator_model)[0]["generated_text"]
 
     return jsonify({'response': response})