Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI | |
| import os | |
| import pymupdf # PyMuPDF | |
| from pptx import Presentation | |
| from sentence_transformers import SentenceTransformer | |
| import torch | |
| from transformers import CLIPProcessor, CLIPModel | |
| from PIL import Image | |
| import chromadb | |
| import numpy as np | |
| from sklearn.decomposition import PCA | |
| app = FastAPI() | |
| # Initialize ChromaDB | |
| client = chromadb.PersistentClient(path="/data/chroma_db") | |
| collection = client.get_or_create_collection(name="knowledge_base") | |
| # File Paths | |
| pdf_file = "Sutures and Suturing techniques.pdf" | |
| pptx_file = "impalnt 1.pptx" | |
| # Initialize Embedding Models | |
| text_model = SentenceTransformer('all-MiniLM-L6-v2') | |
| model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") | |
| processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") | |
| # Image Storage Folder | |
| IMAGE_FOLDER = "/data/extracted_images" | |
| os.makedirs(IMAGE_FOLDER, exist_ok=True) | |
| # Extract Text from PDF | |
| def extract_text_from_pdf(pdf_path): | |
| try: | |
| doc = pymupdf.open(pdf_path) | |
| text = " ".join(page.get_text() for page in doc) | |
| return text.strip() if text else None | |
| except Exception as e: | |
| print(f"Error extracting text from PDF: {e}") | |
| return None | |
| # Extract Text from PPTX | |
| def extract_text_from_pptx(pptx_path): | |
| try: | |
| prs = Presentation(pptx_path) | |
| text = " ".join( | |
| shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text") | |
| ) | |
| return text.strip() if text else None | |
| except Exception as e: | |
| print(f"Error extracting text from PPTX: {e}") | |
| return None | |
| # Extract Images from PDF | |
| def extract_images_from_pdf(pdf_path): | |
| try: | |
| doc = pymupdf.open(pdf_path) | |
| images = [] | |
| for i, page in enumerate(doc): | |
| for img_index, img in enumerate(page.get_images(full=True)): | |
| xref = img[0] | |
| image = doc.extract_image(xref) | |
| img_path = f"{IMAGE_FOLDER}/pdf_image_{i}_{img_index}.{image['ext']}" | |
| with open(img_path, "wb") as f: | |
| f.write(image["image"]) | |
| images.append(img_path) | |
| return images | |
| except Exception as e: | |
| print(f"Error extracting images from PDF: {e}") | |
| return [] | |
| # Extract Images from PPTX | |
| def extract_images_from_pptx(pptx_path): | |
| try: | |
| images = [] | |
| prs = Presentation(pptx_path) | |
| for i, slide in enumerate(prs.slides): | |
| for shape in slide.shapes: | |
| if shape.shape_type == 13: | |
| img_path = f"{IMAGE_FOLDER}/pptx_image_{i}.{shape.image.ext}" | |
| with open(img_path, "wb") as f: | |
| f.write(shape.image.blob) | |
| images.append(img_path) | |
| return images | |
| except Exception as e: | |
| print(f"Error extracting images from PPTX: {e}") | |
| return [] | |
| # Convert Text to Embeddings | |
| def get_text_embedding(text): | |
| return text_model.encode(text).tolist() | |
| # Preload PCA instance globally (to maintain consistency across calls) | |
| pca = PCA(n_components=384) | |
| def get_image_embedding(image_path): | |
| try: | |
| # Load the image | |
| image = Image.open(image_path) | |
| inputs = processor(images=image, return_tensors="pt") | |
| # Extract image embeddings | |
| with torch.no_grad(): | |
| image_embedding = model.get_image_features(**inputs).numpy().flatten() | |
| # Print the actual embedding dimension | |
| print(f"Image embedding shape: {image_embedding.shape}") | |
| """ # CASE 1: Embedding is already 384-dimensional ✅ | |
| if len(image_embedding) == 384: | |
| return image_embedding.tolist() | |
| # CASE 2: Embedding is larger than 384 (e.g., 512) → Apply PCA ✅ | |
| elif len(image_embedding) > 384: | |
| pca = PCA(n_components=384, svd_solver='auto') # Auto solver for stability | |
| image_embedding = pca.fit_transform(image_embedding.reshape(1, -1)).flatten() | |
| print(f"Reduced image embedding shape: {image_embedding.shape}") | |
| # CASE 3: Embedding is smaller than 384 → Apply Padding ❌ | |
| else: | |
| padding = np.zeros(384 - len(image_embedding)) # Create padding vector | |
| image_embedding = np.concatenate((image_embedding, padding)) # Append padding""" | |
| # Truncate to 384 dimensions | |
| image_embedding = image_embedding[:384] | |
| # Print the final embedding shape | |
| print(f"Final Image embedding shape: {image_embedding.shape}") | |
| return image_embedding.tolist() | |
| except Exception as e: | |
| print(f"❌ Error generating image embedding: {e}") | |
| return None | |
| # Store Data in ChromaDB | |
| def store_data(texts, image_paths): | |
| for i, text in enumerate(texts): | |
| if text: | |
| text_embedding = get_text_embedding(text) | |
| if len(text_embedding) == 384: | |
| collection.add(ids=[f"text_{i}"], embeddings=[text_embedding], documents=[text]) | |
| all_embeddings = [get_image_embedding(img_path) for img_path in image_paths if get_image_embedding(img_path) is not None] | |
| if all_embeddings: | |
| all_embeddings = np.array(all_embeddings) | |
| # Apply PCA only if necessary | |
| if all_embeddings.shape[1] != 384: | |
| pca = PCA(n_components=384) | |
| all_embeddings = pca.fit_transform(all_embeddings) | |
| for j, img_path in enumerate(image_paths): | |
| collection.add(ids=[f"image_{j}"], embeddings=[all_embeddings[j].tolist()], documents=[img_path]) | |
| print("Data stored successfully!") | |
| # Process and Store from Files | |
| def process_and_store(pdf_path=None, pptx_path=None): | |
| texts, images = [], [] | |
| if pdf_path: | |
| pdf_text = extract_text_from_pdf(pdf_path) | |
| if pdf_text: | |
| texts.append(pdf_text) | |
| images.extend(extract_images_from_pdf(pdf_path)) | |
| if pptx_path: | |
| pptx_text = extract_text_from_pptx(pptx_path) | |
| if pptx_text: | |
| texts.append(pptx_text) | |
| images.extend(extract_images_from_pptx(pptx_path)) | |
| store_data(texts, images) | |
| # FastAPI Endpoints | |
| def greet_json(): | |
| # Run Data Processing | |
| process_and_store(pdf_path=pdf_file, pptx_path=pptx_file) | |
| return {"Document store": "created!"} | |
| def retrieval(query: str): | |
| try: | |
| query_embedding = get_text_embedding(query) | |
| results = collection.query(query_embeddings=[query_embedding], n_results=5) | |
| #return {"results": results.get("documents", [])} | |
| # Set a similarity threshold (adjust as needed) | |
| SIMILARITY_THRESHOLD = 0.7 | |
| # Extract documents and similarity scores | |
| documents = results.get("documents", [[]])[0] # Ensure we get the first list | |
| distances = results.get("distances", [[]])[0] # Ensure we get the first list | |
| # Filter results based on similarity threshold | |
| filtered_results = [ | |
| doc for doc, score in zip(documents, distances) if score >= SIMILARITY_THRESHOLD | |
| ] | |
| # Return filtered results or indicate no match found | |
| if filtered_results: | |
| return {"results": filtered_results} | |
| else: | |
| return {"results": "No relevant match found in ChromaDB."} | |
| except Exception as e: | |
| return {"error": str(e)} | |
| import pandas as pd | |
| from io import StringIO | |
| import os | |
| import base64 | |
| def save_file_dify(csv_data: str): | |
| # Split into lines | |
| lines = csv_data.split("\n") | |
| # Find the max number of columns | |
| max_cols = max(line.count(",") + 1 for line in lines if line.strip()) | |
| # Normalize all rows to have the same number of columns | |
| fixed_lines = [line + "," * (max_cols - line.count(",") - 1) for line in lines] | |
| # Reconstruct CSV string | |
| fixed_csv_data = "\n".join(fixed_lines) | |
| # Convert CSV string to DataFrame | |
| df = pd.read_csv(StringIO(fixed_csv_data)) | |
| #save in dify dataset and return download link | |
| download_link = get_download_link_dify(df) | |
| return download_link | |
| def get_download_link_dify(df): | |
| # code to save file in dify framework | |
| import requests | |
| # API Configuration | |
| BASE_URL = "http://redmindgpt.redmindtechnologies.com:81/v1" | |
| DATASET_ID = "084ae979-d101-414b-8854-9bbf5d3a442e" | |
| API_KEY = "dataset-feqz5KrqHkFRdWbh2DInt58L" | |
| dataset_name = 'output_dataset' | |
| # Endpoint URL | |
| url = f"{BASE_URL}/datasets/{DATASET_ID}/document/create-by-file" | |
| print(url) | |
| # Headers | |
| headers = { | |
| "Authorization": f"Bearer {API_KEY}" | |
| } | |
| # Data payload (form data as a plain text string) | |
| data_payload = { | |
| "data": """ | |
| { | |
| "indexing_technique": "high_quality", | |
| "process_rule": { | |
| "rules": { | |
| "pre_processing_rules": [ | |
| {"id": "remove_extra_spaces", "enabled": true}, | |
| {"id": "remove_urls_emails", "enabled": true} | |
| ], | |
| "segmentation": { | |
| "separator": "###", | |
| "max_tokens": 500 | |
| } | |
| }, | |
| "mode": "custom" | |
| } | |
| } | |
| """ | |
| } | |
| # Convert DataFrame to binary (in-memory) | |
| file_buffer = dataframe_to_binary(df) | |
| files = { | |
| "file": ("output.xlsx", file_buffer, "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet") | |
| } | |
| # Send the POST request | |
| response = requests.post(url, headers=headers, data=data_payload, files=files) | |
| print(response) | |
| data = response.json() | |
| document_id = data['document']['id'] | |
| # code to get download_url | |
| url = f"http://redmindgpt.redmindtechnologies.com:81/v1/datasets/{DATASET_ID}/documents/{document_id}/upload-file" | |
| response = requests.get(url, headers=headers) | |
| print(response) | |
| download_url = response.json().get("download_url") | |
| download_url = download_url.replace("download/","") | |
| return download_url | |
| def dataframe_to_binary(df): | |
| import io | |
| # Create a BytesIO stream | |
| output = io.BytesIO() | |
| # Write the DataFrame to this in-memory buffer as an Excel file | |
| df.to_excel(output, index=False, engine="openpyxl") | |
| # Move the cursor to the beginning of the stream | |
| output.seek(0) | |
| return output | |