Vision All You Need by Softlandia

├── vrag ├── __init__.py ├── app.py ├── pdf_to_image.py ├── qdrant_client.py ├── vrag.py └── colpali.py ├── frontend ├── .eslintignore ├── src │ ├── vite-env.d.ts │ ├── App.css.ts │ ├── Question │ │ ├── Question.css.ts │ │ └── Question.tsx │ ├── theme.ts │ ├── main.tsx │ ├── ColorSchemeToggle │ │ ├── ColorSchemeToggle.css.ts │ │ └── ColorSchemeToggle.tsx │ ├── styles │ │ └── global.css.ts │ ├── Header │ │ ├── Header.css.ts │ │ └── Header.tsx │ ├── Dropzone │ │ ├── Dropzone.css.ts │ │ └── Dropzone.tsx │ ├── AttentionMap │ │ └── AttentionMap.tsx │ ├── api │ │ └── index.ts │ └── App.tsx ├── public │ ├── favicon.ico │ ├── softlandia_logo_h_white_1.png │ └── Softlandia_logo_Hor_color_small.png ├── README.md ├── tsconfig.node.json ├── vite.config.ts ├── postcss.config.cjs ├── .eslintrc.cjs ├── tsconfig.json ├── index.html ├── package.json └── .gitignore ├── requirements.txt ├── LICENSE ├── README.md ├── .gitignore └── main.py /vrag/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /frontend/.eslintignore: -------------------------------------------------------------------------------- 1 | *.js 2 | *.cjs 3 | *.mjs -------------------------------------------------------------------------------- /frontend/src/vite-env.d.ts: -------------------------------------------------------------------------------- 1 | /// 2 | -------------------------------------------------------------------------------- /vrag/app.py: -------------------------------------------------------------------------------- 1 | from modal import App 2 | 3 | app = App(name="vision-is-all-you-need") 4 | -------------------------------------------------------------------------------- /frontend/public/favicon.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Softlandia-Ltd/vision-is-all-you-need/HEAD/frontend/public/favicon.ico -------------------------------------------------------------------------------- /frontend/src/App.css.ts: -------------------------------------------------------------------------------- 1 | import { style } from "@vanilla-extract/css"; 2 | 3 | export const sources = style({ 4 | overflow: "auto", 5 | }); -------------------------------------------------------------------------------- /frontend/public/softlandia_logo_h_white_1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Softlandia-Ltd/vision-is-all-you-need/HEAD/frontend/public/softlandia_logo_h_white_1.png -------------------------------------------------------------------------------- /frontend/src/Question/Question.css.ts: -------------------------------------------------------------------------------- 1 | import { style } from "@vanilla-extract/css"; 2 | 3 | export const questionInput = style({ 4 | fontSize: "1.5rem", 5 | }); 6 | -------------------------------------------------------------------------------- /frontend/public/Softlandia_logo_Hor_color_small.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Softlandia-Ltd/vision-is-all-you-need/HEAD/frontend/public/Softlandia_logo_Hor_color_small.png -------------------------------------------------------------------------------- /frontend/src/theme.ts: -------------------------------------------------------------------------------- 1 | import { createTheme } from "@mantine/core"; 2 | import { themeToVars } from "@mantine/vanilla-extract"; 3 | 4 | export const theme = createTheme({}); 5 | export const vars = themeToVars(theme); 6 | -------------------------------------------------------------------------------- /frontend/README.md: -------------------------------------------------------------------------------- 1 | # Mantine + Vite + Vanilla extract template 2 | 3 | Get started with the template by clicking `Use this template` button on the top of the page. 4 | 5 | [Documentation](https://mantine.dev/styles/vanilla-extract/) 6 | -------------------------------------------------------------------------------- /frontend/tsconfig.node.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "composite": true, 4 | "skipLibCheck": true, 5 | "module": "ESNext", 6 | "moduleResolution": "bundler", 7 | "allowSyntheticDefaultImports": true 8 | }, 9 | "include": ["vite.config.ts"] 10 | } 11 | -------------------------------------------------------------------------------- /frontend/src/main.tsx: -------------------------------------------------------------------------------- 1 | import React from "react"; 2 | import ReactDOM from "react-dom/client"; 3 | import App from "./App.tsx"; 4 | 5 | ReactDOM.createRoot(document.getElementById("root") as HTMLElement).render( 6 | 7 | 8 | 9 | ); 10 | -------------------------------------------------------------------------------- /frontend/vite.config.ts: -------------------------------------------------------------------------------- 1 | import { defineConfig } from "vite"; 2 | import { vanillaExtractPlugin } from "@vanilla-extract/vite-plugin"; 3 | import react from "@vitejs/plugin-react"; 4 | 5 | // https://vitejs.dev/config/ 6 | export default defineConfig({ 7 | plugins: [react(), vanillaExtractPlugin()], 8 | }); 9 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | pydantic==2.9.1 2 | qdrant-client==1.11.1 3 | pypdfium2==4.30.0 4 | fastapi[standard]==0.114.0 5 | python-dotenv==1.0.1 6 | colpali_engine==0.3.1 7 | openai==1.44.1 8 | opencv_python_headless==4.10.0.84 9 | tqdm==4.66.5 10 | transformers>=4.45.0 11 | sse-starlette==2.1.3 12 | vidore_benchmark==4.0.1 13 | einops==0.8.0 14 | torch 15 | -------------------------------------------------------------------------------- /frontend/postcss.config.cjs: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | plugins: { 3 | "postcss-preset-mantine": {}, 4 | "postcss-simple-vars": { 5 | variables: { 6 | "mantine-breakpoint-xs": "36em", 7 | "mantine-breakpoint-sm": "48em", 8 | "mantine-breakpoint-md": "62em", 9 | "mantine-breakpoint-lg": "75em", 10 | "mantine-breakpoint-xl": "88em", 11 | }, 12 | }, 13 | }, 14 | }; 15 | -------------------------------------------------------------------------------- /frontend/.eslintrc.cjs: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | env: { browser: true, es2020: true }, 3 | extends: [ 4 | 'eslint:recommended', 5 | 'plugin:@typescript-eslint/recommended', 6 | 'plugin:react-hooks/recommended', 7 | ], 8 | parser: '@typescript-eslint/parser', 9 | parserOptions: { ecmaVersion: 'latest', sourceType: 'module' }, 10 | plugins: ['react-refresh'], 11 | rules: { 12 | 'react-refresh/only-export-components': 'warn', 13 | }, 14 | } 15 | -------------------------------------------------------------------------------- /frontend/src/ColorSchemeToggle/ColorSchemeToggle.css.ts: -------------------------------------------------------------------------------- 1 | import { style } from "@vanilla-extract/css"; 2 | import { vars } from "./../theme"; 3 | 4 | export const icon = style({ 5 | width: "22px", 6 | height: "22px", 7 | }); 8 | 9 | export const dark = style({ 10 | selectors: { 11 | [vars.darkSelector]: { 12 | display: "none", 13 | }, 14 | [vars.lightSelector]: { 15 | display: "block", 16 | }, 17 | }, 18 | }); 19 | 20 | export const light = style({ 21 | selectors: { 22 | [vars.lightSelector]: { 23 | display: "none", 24 | }, 25 | [vars.darkSelector]: { 26 | display: "block", 27 | }, 28 | }, 29 | }); 30 | -------------------------------------------------------------------------------- /frontend/tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "ESNext", 4 | "lib": ["DOM", "DOM.Iterable", "ESNext"], 5 | "module": "ESNext", 6 | "skipLibCheck": true, 7 | 8 | /* Bundler mode */ 9 | "moduleResolution": "bundler", 10 | "allowImportingTsExtensions": true, 11 | "resolveJsonModule": true, 12 | "isolatedModules": true, 13 | "noEmit": true, 14 | "jsx": "react-jsx", 15 | 16 | /* Linting */ 17 | "strict": true, 18 | "noUnusedLocals": true, 19 | "noUnusedParameters": true, 20 | "noFallthroughCasesInSwitch": true 21 | }, 22 | "include": ["src"], 23 | "references": [{ "path": "./tsconfig.node.json" }] 24 | } 25 | -------------------------------------------------------------------------------- /vrag/pdf_to_image.py: -------------------------------------------------------------------------------- 1 | import pypdfium2 as pdfium 2 | import numpy as np 3 | 4 | 5 | def images_from_pdf_bytes(pdf_bytes: bytes) -> list[np.ndarray]: 6 | pdf = pdfium.PdfDocument(pdf_bytes) 7 | return images_from_document(pdf) 8 | 9 | 10 | def images_from_pdf_path(pdf_path: str) -> list[np.ndarray]: 11 | pdf = pdfium.PdfDocument(pdf_path) 12 | return images_from_document(pdf) 13 | 14 | 15 | def images_from_document(document: pdfium.PdfDocument) -> list[np.ndarray]: 16 | images = [] 17 | scale = 92 / 72 18 | 19 | for i in range(len(document)): 20 | page = document[i] 21 | img = page.render(scale=scale, rev_byteorder=True).to_numpy() # type: ignore 22 | images.append(img) 23 | 24 | return images 25 | -------------------------------------------------------------------------------- /frontend/index.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 10 | Vision All You Need by Softlandia 11 | 12 | 13 | 17 | 18 | 19 |

20 | 21 | 22 | 23 | -------------------------------------------------------------------------------- /frontend/src/ColorSchemeToggle/ColorSchemeToggle.tsx: -------------------------------------------------------------------------------- 1 | import { ActionIcon, Tooltip, useMantineColorScheme } from "@mantine/core"; 2 | import { IconSun, IconMoon } from "@tabler/icons-react"; 3 | import cx from "clsx"; 4 | import * as classes from "./ColorSchemeToggle.css"; 5 | 6 | export function ColorSchemeToggle() { 7 | const { setColorScheme, colorScheme } = useMantineColorScheme(); 8 | 9 | return ( 10 | 13 | 15 | setColorScheme(colorScheme === "light" ? "dark" : "light") 16 | } 17 | variant="default" 18 | size="lg" 19 | aria-label="Toggle color scheme" 20 | > 21 | 22 | 23 | 24 | 25 | ); 26 | } 27 | -------------------------------------------------------------------------------- /frontend/src/styles/global.css.ts: -------------------------------------------------------------------------------- 1 | import { globalStyle } from "@vanilla-extract/css"; 2 | 3 | // Default style resets 4 | 5 | globalStyle("html, body", { 6 | margin: 0, 7 | padding: 0, 8 | minHeight: "100vh", 9 | height: "100%", 10 | fontFamily: "'Quicksand', sans-serif", 11 | }); 12 | 13 | globalStyle("h1, h2, h3, h4, h5, h6, p, a, span, label, div", { 14 | fontWeight: "inherit", 15 | fontSize: "inherit", 16 | color: "inherit", 17 | textDecoration: "none", 18 | margin: 0, 19 | padding: 0, 20 | fontFamily: "'Quicksand', sans-serif", 21 | }); 22 | 23 | globalStyle("button", { 24 | WebkitTouchCallout: "none", // iOS Safari 25 | WebkitUserSelect: "none", // Safari 26 | MozUserSelect: "none", // Old versions of Firefox 27 | msUserSelect: "none", // Internet Explorer / Edge 28 | userSelect: "none", // Chrome, Edge, Opera and Firefox 29 | }); 30 | 31 | // Root styles 32 | 33 | globalStyle("#root", { 34 | height: "100%", 35 | fontFamily: "'Quicksand', sans-serif", 36 | }); 37 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Softlandia 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /frontend/src/Header/Header.css.ts: -------------------------------------------------------------------------------- 1 | import { style } from "@vanilla-extract/css"; 2 | import { vars } from "../theme"; 3 | import { rem } from "@mantine/core"; 4 | 5 | export const title = style({ 6 | color: vars.colors.black, 7 | fontSize: rem(50), 8 | fontWeight: 900, 9 | letterSpacing: rem(-2), 10 | selectors: { 11 | [vars.darkSelector]: { 12 | color: vars.colors.white, 13 | }, 14 | }, 15 | 16 | "@media": { 17 | [vars.smallerThan("md")]: { 18 | fontSize: rem(25), 19 | }, 20 | [vars.smallerThan("sm")]: { 21 | fontSize: rem(20), 22 | }, 23 | }, 24 | }); 25 | 26 | export const header = style({ 27 | height: 60, 28 | backgroundColor: vars.colors.body, 29 | borderBottom: `1px solid ${vars.colors.gray[3]}`, 30 | borderBottomColor: vars.colors.gray[3], 31 | display: "flex", 32 | justifyContent: "center", 33 | alignItems: "center", 34 | paddingBottom: 5, 35 | selectors: { 36 | [vars.darkSelector]: { 37 | borderBottomColor: vars.colors.dark[4], 38 | }, 39 | }, 40 | "@media": { 41 | [vars.smallerThan("md")]: { 42 | height: 90, 43 | paddingBottom: 10, 44 | }, 45 | [vars.smallerThan("sm")]: { 46 | height: 90, 47 | paddingBottom: 10, 48 | }, 49 | }, 50 | }); 51 | -------------------------------------------------------------------------------- /frontend/src/Question/Question.tsx: -------------------------------------------------------------------------------- 1 | import { Textarea, ActionIcon, useMantineTheme, rem } from "@mantine/core"; 2 | import { IconArrowRight } from "@tabler/icons-react"; 3 | import { useState } from "react"; 4 | import * as classes from "./Question.css"; 5 | 6 | 7 | export interface QuestionProps { 8 | onSubmit: (question: string) => void; 9 | loading: boolean; 10 | disabled: boolean; 11 | } 12 | 13 | export function Question(props: QuestionProps) { 14 | const theme = useMantineTheme(); 15 | const [question, setQuestion] = useState(""); 16 | 17 | return ( 18 |

setQuestion(event.currentTarget.value)}
25 |       rightSectionWidth={82}
26 |       classNames={{
27 |         input: classes.questionInput
28 |       }}
29 |       rightSection={
30 |         <ActionIcon
31 |           size={62}
32 |           radius="md"
33 |           disabled={props.loading || props.disabled}
34 |           loading={props.loading}
35 |           color={theme.primaryColor}
36 |           variant="filled"
37 |           onClick={() => props.onSubmit(question)}
38 |         >
39 |           <IconArrowRight
40 |             style={{ width: rem(18), height: rem(18) }}
41 |             stroke={1.5}
42 |           />
43 |         </ActionIcon>
44 |       }
45 |     />
46 |   );
47 | }
48 |

--------------------------------------------------------------------------------
/frontend/package.json:
--------------------------------------------------------------------------------
 1 | {
 2 |   "name": "vision-is-all-you-need",
 3 |   "private": true,
 4 |   "version": "0.0.1",
 5 |   "type": "module",
 6 |   "scripts": {
 7 |     "dev": "vite dev",
 8 |     "build": "tsc && vite build",
 9 |     "lint": "eslint src --ext ts,tsx --report-unused-disable-directives --max-warnings 0",
10 |     "preview": "vite preview"
11 |   },
12 |   "dependencies": {
13 |     "@mantine/core": "7.12.2",
14 |     "@mantine/dropzone": "^7.12.2",
15 |     "@mantine/hooks": "7.12.2",
16 |     "@mantine/vanilla-extract": "7.12.2",
17 |     "@microsoft/fetch-event-source": "^2.0.1",
18 |     "@tabler/icons-react": "^3.17.0",
19 |     "@types/react-syntax-highlighter": "^15.5.13",
20 |     "@vanilla-extract/css": "^1.15.3",
21 |     "react": "^18.3.1",
22 |     "react-dom": "^18.3.1",
23 |     "react-markdown": "^8.0.6",
24 |     "react-syntax-highlighter": "^15.5.0",
25 |     "remark-gfm": "^4.0.0"
26 |   },
27 |   "devDependencies": {
28 |     "@types/react": "^18.3.3",
29 |     "@types/react-dom": "^18.3.0",
30 |     "@typescript-eslint/eslint-plugin": "^7.13.1",
31 |     "@typescript-eslint/parser": "^7.13.1",
32 |     "@vanilla-extract/vite-plugin": "^4.0.11",
33 |     "@vitejs/plugin-react": "^4.3.1",
34 |     "eslint": "^8.57.0",
35 |     "eslint-plugin-react-hooks": "^4.6.2",
36 |     "eslint-plugin-react-refresh": "^0.4.7",
37 |     "postcss": "^8.4.38",
38 |     "postcss-preset-mantine": "1.17.0",
39 |     "postcss-simple-vars": "^7.0.1",
40 |     "typescript": "^5.5.2",
41 |     "vite": "^5.3.1"
42 |   }
43 | }
44 |

--------------------------------------------------------------------------------
/frontend/src/Dropzone/Dropzone.css.ts:
--------------------------------------------------------------------------------
 1 | import { style } from "@vanilla-extract/css";
 2 | import { vars } from "../theme";
 3 | 
 4 | export const wrapper = style({
 5 |   display: "flex",
 6 |   flexDirection: "column",
 7 |   alignItems: "center",
 8 |   gap: "16px",
 9 |   width: "100%",
10 |   maxWidth: "100%",
11 | });
12 | 
13 | export const dropzone = style({
14 |   width: "100%",
15 |   padding: "20px",
16 |   border: `2px dashed ${vars.colors.gray[4]}`,
17 |   borderRadius: vars.radius.md,
18 |   transition: "background-color 0.3s ease",
19 |   backgroundColor: vars.colors.gray[1],
20 | 
21 |   selectors: {
22 |     "&:hover": {
23 |       backgroundColor: vars.colors.gray[2],
24 |     },
25 |     [vars.darkSelector]: {
26 |       border: `2px dashed ${vars.colors.dark[4]}`,
27 |       backgroundColor: vars.colors.dark[7],
28 |     },
29 |     [vars.darkSelector + " &:hover"]: {
30 |       backgroundColor: vars.colors.dark[6],
31 |     },
32 |   },
33 | });
34 | 
35 | export const control = style({
36 |   width: "200px",
37 |   display: "block",
38 |   margin: "0 auto",
39 |   textAlign: "center",
40 |   padding: "12px 24px",
41 |   borderRadius: vars.radius.xl,
42 |   backgroundColor: vars.colors.blue[6],
43 |   color: vars.colors.white,
44 |   fontSize: "16px",
45 |   fontWeight: 700,
46 | 
47 |   selectors: {
48 |     "&:hover": {
49 |       backgroundColor: vars.colors.blue[7],
50 |     },
51 |     [vars.darkSelector]: {
52 |       backgroundColor: vars.colors.blue[8],
53 |     },
54 |     [vars.darkSelector + " &:hover"]: {
55 |       backgroundColor: vars.colors.blue[9],
56 |     },
57 |   },
58 | });
59 |

--------------------------------------------------------------------------------
/frontend/src/AttentionMap/AttentionMap.tsx:
--------------------------------------------------------------------------------
 1 | import { useLayoutEffect, useRef } from "react";
 2 | 
 3 | export interface AttentionMapProps {
 4 |   tokenMap: number[][] | null;
 5 |   width: number;
 6 |   height: number;
 7 |   color: "red" | "blue" | "green";
 8 |   opacity: number;
 9 | }
10 | 
11 | export function AttentionMap({
12 |   tokenMap,
13 |   width,
14 |   height,
15 |   color,
16 |   opacity,
17 | }: AttentionMapProps) {
18 |   const canvasRef = useRef<HTMLCanvasElement>(null);
19 | 
20 |   const drawHeatmap = (
21 |     map: number[][] | null,
22 |     ctx: CanvasRenderingContext2D
23 |   ) => {
24 |     if (!map || !ctx) return;
25 | 
26 |     const canvasWidth = width;
27 |     const canvasHeight = height;
28 |     ctx.clearRect(0, 0, canvasWidth, canvasHeight);
29 | 
30 |     const mapHeight = map.length;
31 |     const mapWidth = map[0].length;
32 |     const cellWidth = canvasWidth / mapWidth;
33 |     const cellHeight = canvasHeight / mapHeight;
34 | 
35 |     for (let row = 0; row < mapHeight; row++) {
36 |       for (let col = 0; col < mapWidth; col++) {
37 |         const attentionValue = map[row][col];
38 | 
39 |         const intensity = Math.min(Math.max(attentionValue * 255, 0), 255);
40 |         const red = color === "red" ? 255 : 0;
41 |         const green = color === "green" ? 255 : 0;
42 |         const blue = color === "blue" ? 255 : 0;
43 |         ctx.fillStyle = `rgba(${red}, ${green}, ${blue}, ${
44 |           (intensity / 255) * (opacity / 100)
45 |         })`;
46 | 
47 |         ctx.fillRect(col * cellWidth, row * cellHeight, cellWidth, cellHeight);
48 |       }
49 |     }
50 |   };
51 | 
52 |   useLayoutEffect(() => {
53 |     const canvas = canvasRef.current;
54 |     if (!canvas) return;
55 |     const ctx = canvas.getContext("2d");
56 |     if (ctx && tokenMap) {
57 |       drawHeatmap(tokenMap, ctx);
58 |     }
59 |   }, [tokenMap, color, opacity]);
60 | 
61 |   return (
62 |     <canvas
63 |       ref={canvasRef}
64 |       width={width}
65 |       height={height}
66 |       style={{
67 |         position: "absolute",
68 |         top: 0,
69 |         left: 0,
70 |         pointerEvents: "auto",
71 |       }}
72 |     />
73 |   );
74 | }
75 |

--------------------------------------------------------------------------------
/vrag/qdrant_client.py:
--------------------------------------------------------------------------------
 1 | import uuid
 2 | from pydantic import BaseModel
 3 | from qdrant_client import AsyncQdrantClient, models
 4 | from typing import Any
 5 | 
 6 | 
 7 | class Result(BaseModel):
 8 |     id: int | str
 9 |     score: float
10 |     payload: dict[str, Any]
11 | 
12 | 
13 | class InMemoryQdrant:
14 |     def __init__(self) -> None:
15 |         self.client = AsyncQdrantClient(":memory:")
16 | 
17 |     async def does_collection_exist(self, name: str) -> bool:
18 |         return await self.client.collection_exists(name)
19 | 
20 |     async def create_collection(self, name: str):
21 |         exists = await self.client.collection_exists(name)
22 | 
23 |         if exists:
24 |             return
25 | 
26 |         return await self.client.create_collection(
27 |             collection_name=name,
28 |             vectors_config=models.VectorParams(
29 |                 size=128,
30 |                 distance=models.Distance.COSINE,
31 |                 multivector_config=models.MultiVectorConfig(
32 |                     comparator=models.MultiVectorComparator.MAX_SIM
33 |                 ),
34 |             ),
35 |         )
36 | 
37 |     async def search_collection(
38 |         self, collection: str, query_vector: list[list[float]], count: int
39 |     ) -> list[Result]:
40 |         result = await self.client.query_points(
41 |             collection, query=query_vector, limit=count
42 |         )
43 |         return [
44 |             Result(id=point.id, score=point.score, payload=point.payload or {})
45 |             for point in result.points
46 |         ]
47 | 
48 |     async def upsert_points(
49 |         self,
50 |         collection: str,
51 |         pdf_name: str,
52 |         embeddings: list[list[list[float]]],
53 |         encoded_images: list[str],
54 |     ) -> models.UpdateResult:
55 |         points = [
56 |             models.PointStruct(
57 |                 id=str(uuid.uuid4()),
58 |                 vector=embedding,
59 |                 payload={
60 |                     "page": idx + 1,
61 |                     "name": pdf_name,
62 |                     "image": encoded_images[idx],
63 |                 },
64 |             )
65 |             for idx, embedding in enumerate(embeddings)
66 |         ]
67 |         return await self.client.upsert(collection, points)
68 |

--------------------------------------------------------------------------------
/frontend/src/Dropzone/Dropzone.tsx:
--------------------------------------------------------------------------------
 1 | import { useRef } from "react";
 2 | import { Text, Group, rem, useMantineTheme } from "@mantine/core";
 3 | import { Dropzone, MIME_TYPES } from "@mantine/dropzone";
 4 | import { IconCloudUpload, IconX, IconDownload } from "@tabler/icons-react";
 5 | import * as classes from "./Dropzone.css";
 6 | 
 7 | export interface DropzoneBoxProps {
 8 |   uploading: boolean;
 9 |   onSubmit: (files: File[]) => void;
10 | }
11 | 
12 | export function DropzoneBox(props: DropzoneBoxProps) {
13 |   const theme = useMantineTheme();
14 |   const openRef = useRef<() => void>(null);
15 | 
16 |   return (
17 |     <div className={classes.wrapper}>
18 |       <Dropzone
19 |         disabled={props.uploading}
20 |         openRef={openRef}
21 |         onDrop={(files) => props.onSubmit(files)}
22 |         className={classes.dropzone}
23 |         radius="md"
24 |         accept={[MIME_TYPES.pdf]}
25 |         maxSize={30 * 1024 ** 2}
26 |       >
27 |         <div style={{ pointerEvents: "none" }}>
28 |           <Group justify="center">
29 |             <Dropzone.Accept>
30 |               <IconDownload
31 |                 style={{ width: rem(50), height: rem(50) }}
32 |                 color={theme.colors.blue[6]}
33 |                 stroke={1.5}
34 |               />
35 |             </Dropzone.Accept>
36 |             <Dropzone.Reject>
37 |               <IconX
38 |                 style={{ width: rem(50), height: rem(50) }}
39 |                 color={theme.colors.red[6]}
40 |                 stroke={1.5}
41 |               />
42 |             </Dropzone.Reject>
43 |             <Dropzone.Idle>
44 |               <IconCloudUpload
45 |                 style={{ width: rem(50), height: rem(50) }}
46 |                 stroke={1.5}
47 |               />
48 |             </Dropzone.Idle>
49 |           </Group>
50 | 
51 |           <Text ta="center" fw={700} fz="lg">
52 |             <Dropzone.Accept>Drop files here</Dropzone.Accept>
53 |             <Dropzone.Reject>Pdf file less than 30mb</Dropzone.Reject>
54 |             <Dropzone.Idle>Upload PDFs</Dropzone.Idle>
55 |           </Text>
56 |           <Text ta="center" fz="sm" c="dimmed">
57 |             Only <i>.pdf</i> files that are less than 30mb in size are accepted.
58 |             There's no persistent storage, the index is built on the fly and
59 |             will be lost when the underlying container dies. <br /> If the
60 |             upload takes long, it's likely that the container has been
61 |             terminated and you need to wait for a little while for the new
62 |             container to be up and running.
63 |           </Text>
64 |         </div>
65 |       </Dropzone>
66 |     </div>
67 |   );
68 | }
69 |

--------------------------------------------------------------------------------
/frontend/src/Header/Header.tsx:
--------------------------------------------------------------------------------
 1 | import {
 2 |   Title,
 3 |   Text,
 4 |   Group,
 5 |   Image,
 6 |   useMantineColorScheme,
 7 |   ActionIcon,
 8 |   Tooltip,
 9 |   Flex,
10 | } from "@mantine/core";
11 | import * as classes from "./Header.css";
12 | import { ColorSchemeToggle } from "../ColorSchemeToggle/ColorSchemeToggle";
13 | import { IconBrandGithub } from "@tabler/icons-react";
14 | import { IconArticle } from "@tabler/icons-react";
15 | 
16 | export default function Header() {
17 |   const { colorScheme } = useMantineColorScheme();
18 |   return (
19 |     <header className={classes.header}>
20 |       <Group w="100%" gap="lg">
21 |         <Flex
22 |           direction={{ base: "column", sm: "row" }}
23 |           align="center"
24 |           gap={{ base: 5, sm: "none" }}
25 |         >
26 |           <Group w="100%" gap="lg" pr={{ base: "none", sm: "md" }}>
27 |             <Title className={classes.title} ta="center">
28 |               <Text
29 |                 inherit
30 |                 variant="gradient"
31 |                 component="span"
32 |                 gradient={{ from: "pink", to: "yellow" }}
33 |               >
34 |                 Vision is All You Need
35 |               </Text>
36 |             </Title>
37 |           </Group>
38 |           <Group wrap="nowrap">
39 |             <Title ta="center">
40 |               <Text
41 |                 fz={{ base: "sm", sm: "lg" }}
42 |                 c="dimmed"
43 |                 variant="gradient"
44 |                 gradient={{ from: "pink", to: "yellow" }}
45 |               >
46 |                 by
47 |               </Text>
48 |             </Title>
49 |             <a href="https://softlandia.fi">
50 |               {colorScheme === "dark" ? (
51 |                 <Image
52 |                   h={{ base: 20, sm: 30 }}
53 |                   w="auto"
54 |                   fit="contain"
55 |                   src={"/softlandia_logo_h_white_1.png"}
56 |                 />
57 |               ) : (
58 |                 <Image
59 |                   h={{ base: 20, sm: 30 }}
60 |                   w="auto"
61 |                   fit="contain"
62 |                   src={"/Softlandia_logo_Hor_color_small.png"}
63 |                 />
64 |               )}
65 |             </a>
66 |           </Group>
67 |         </Flex>
68 |         <Group ml={"auto"}>
69 |           <Tooltip label="Background blog post">
70 |             <ActionIcon
71 |               size={"lg"}
72 |               component="a"
73 |               variant="default"
74 |               aria-label="Background blog post"
75 |               href="https://softlandia.fi/en/blog/building-a-rag-tired-of-chunking-maybe-vision-is-all-you-need"
76 |             >
77 |               <IconArticle />
78 |             </ActionIcon>
79 |           </Tooltip>
80 |           <Tooltip label="GitHub repository">
81 |             <ActionIcon
82 |               size={"lg"}
83 |               component="a"
84 |               variant="default"
85 |               aria-label="GitHub repository"
86 |               href="https://github.com/Softlandia-Ltd/vision-is-all-you-need"
87 |             >
88 |               <IconBrandGithub />
89 |             </ActionIcon>
90 |           </Tooltip>
91 |           <ColorSchemeToggle />
92 |         </Group>
93 |       </Group>
94 |     </header>
95 |   );
96 | }
97 |

--------------------------------------------------------------------------------
/frontend/src/api/index.ts:
--------------------------------------------------------------------------------
  1 | import { fetchEventSource } from "@microsoft/fetch-event-source";
  2 | const backendUrl = import.meta.env.VITE_BACKEND_URL;
  3 | 
  4 | type fetcherArgs<T> = {
  5 |   method: string;
  6 |   endpoint: string;
  7 |   body?: T | FormData;
  8 |   stream?: boolean;
  9 | };
 10 | 
 11 | type eventSourceArgs<T> = {
 12 |   onMessage: (data: T, event: string) => void;
 13 |   onOpen?: () => void;
 14 |   onClose?: () => void;
 15 |   onError?: (err: Error) => void;
 16 | };
 17 | 
 18 | type eventSourceFetcherArgs<C, T> = fetcherArgs<C> & eventSourceArgs<T>;
 19 | 
 20 | const initRequest = (args: fetcherArgs<any>) => {
 21 |   const { method, endpoint, body } = args;
 22 | 
 23 |   let headers: HeadersInit = {
 24 |     accept: "application/json",
 25 |   };
 26 | 
 27 |   if (!(body instanceof FormData)) {
 28 |     headers["Content-Type"] = "application/json";
 29 |   }
 30 | 
 31 |   const url = backendUrl + endpoint;
 32 | 
 33 |   const request = new Request(url, {
 34 |     method,
 35 |     headers,
 36 |     body: body && !(body instanceof FormData) ? JSON.stringify(body) : body,
 37 |   });
 38 | 
 39 |   return request;
 40 | };
 41 | 
 42 | const fetchStream = async <C, T>(args: eventSourceFetcherArgs<C, T>) => {
 43 |   const request = initRequest(args);
 44 |   const headers: Record<string, string> = {};
 45 |   request.headers.forEach((value, key) => {
 46 |     headers[key] = value;
 47 |   });
 48 | 
 49 |   const res = await fetchEventSource(request, {
 50 |     headers,
 51 |     openWhenHidden: true,
 52 |     async onopen(response) {
 53 |       if (response.ok && response.status === 200) {
 54 |         args.onOpen && args.onOpen();
 55 |       } else if (
 56 |         response.status >= 400 &&
 57 |         response.status < 500 &&
 58 |         response.status !== 429
 59 |       ) {
 60 |         console.error("Error connecting to server", response.statusText);
 61 |         args.onError && args.onError(new Error(response.statusText));
 62 |       }
 63 |     },
 64 |     onmessage(msg) {
 65 |       console.log(msg);
 66 |       try {
 67 |         args.onMessage(JSON.parse(msg.data) as T, msg.event);
 68 |       } catch {
 69 |         args.onMessage(msg.data as T, msg.event);
 70 |       }
 71 |       if (msg.event === "complete") {
 72 |         console.log("Stream complete");
 73 |       }
 74 |     },
 75 |     onclose() {
 76 |       console.log("Connection closed by the server");
 77 |       args.onClose && args.onClose();
 78 |     },
 79 |     onerror(err) {
 80 |       console.log("There was an error from server", err);
 81 |       args.onError && args.onError(err);
 82 |     },
 83 |   });
 84 | 
 85 |   return res;
 86 | };
 87 | 
 88 | export const postStream = async <C, T>(
 89 |   endpoint: string,
 90 |   body: C,
 91 |   onMessage: (data: T, event: string) => void,
 92 |   onOpen?: () => void,
 93 |   onClose?: () => void,
 94 |   onError?: (err: Error) => void
 95 | ) =>
 96 |   await fetchStream<C, T>({
 97 |     method: "POST",
 98 |     endpoint,
 99 |     body,
100 |     onMessage,
101 |     onOpen,
102 |     onClose,
103 |     onError,
104 |   });
105 | 
106 | export const postFilesStream = async <T>(
107 |   endpoint: string,
108 |   body: FormData,
109 |   onMessage: (data: T, event: string) => void,
110 |   onOpen?: () => void,
111 |   onClose?: () => void,
112 |   onError?: (err: Error) => void
113 | ) => {
114 |   await fetchStream<FormData, T>({
115 |     method: "POST",
116 |     endpoint,
117 |     body,
118 |     onMessage,
119 |     onOpen,
120 |     onClose,
121 |     onError,
122 |   });
123 | };
124 |

--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Vision is All You Need: V-RAG (Vision RAG) Demo
 2 | 
 3 | This is a demo of the Vision RAG (V-RAG) architecture.
 4 | 
 5 | The V-RAG architecture utilizes a vision language model (VLM) to embed pages of PDF files (or any other document) as vectors directly, without the tedious chunking process.
 6 | 
 7 | https://github.com/user-attachments/assets/44d72558-80e8-498d-903a-88672ec6ebd7
 8 | 
 9 | ![Screenshot 2024-10-15 145540](https://github.com/user-attachments/assets/ed834b36-843f-43a4-876b-a30b0c1b88af)
10 | 
11 | Check out the background blog post: https://softlandia.fi/en/blog/building-a-rag-tired-of-chunking-maybe-vision-is-all-you-need
12 | 
13 | # How does V-RAG work?
14 | 
15 | 1. The pages of a PDF file are converted to images.
16 |    - In theory these images can be anything, but the current demo uses PDF files since the underlying model has been trained on PDF files
17 |    - `pypdfium` is used to convert the PDF pages to images
18 | 2. The images are passed through a VLM to get the embeddings.
19 |    - ColPali is used as the VLM in this demo
20 | 3. The embeddings are stored in a database
21 |    - QDrant is used as the vector database in this demo
22 | 4. The user passes a query to the V-RAG system
23 | 5. The query is passed through the VLM to get the query embedding
24 | 6. The query embedding is used to search the vector database for similar embeddings
25 | 7. The user query and images of the best matches from the search are passed again to a model that can understand images
26 |    - we use GPT4o or GPT4o-mini in this demo
27 | 8. The model generates a response based on the query and the images
28 | 
29 | # How to run the demo?
30 | 
31 | Make sure tou have an account in Hugging Face. Make sure you are logged into Hugging Face using `transformers-cli login`.
32 | 
33 | For OpenAI API, you need to have an API key. You can get it from here: https://platform.openai.com/account/api-keys
34 | 
35 | You can place the keys to the dotenv file:
36 | 
37 | ```
38 | OPENAI_API_KEY=
39 | HF_TOKEN=
40 | ```
41 | 
42 | Then, you can run the demo by following these steps:
43 | 
44 | 1. Install Python 3.11 or higher
45 | 2. `pip install modal`
46 | 3. `modal setup`
47 | 4. `modal serve \main.py`
48 | 
49 | # How to use the demo from the provided API?
50 | 
51 | 1. Open your browser and go to the url provided by Modal and append `/docs` to the url
52 | 2. Click on the `POST /collections` endpoint
53 | 3. Click on the `Try it out` button
54 | 4. Upload a PDF file
55 | 5. Click on the `Execute` button
56 | 
57 | This will index the PDF file in to in-memory vector database. This will take some time depending on the size of the PDF file and the GPU you are using in Modal. The current demo is using a A10G GPU.
58 | 
59 | You can now search for similar pages using the `POST /search` endpoint.
60 | 
61 | The endpoint sends the page images and the query to the OpenAI API and returns the response. 
62 | 
63 | # Frontend
64 | 
65 | You can also use the frontend to interact with the API. To setup the frontend for local development, follow these steps:
66 | 
67 | 1. Install Node.js
68 | 2. `cd frontend`
69 |    - modify you `.env.development` file and add your `VITE_BACKEND_URL`
70 | 3. `npm install`
71 | 4. `npm run dev`
72 | 
73 | This will start the frontend on `http://localhost:5173`
74 | 
75 | # How to deploy the demo?
76 | 
77 | You can deploy the demo to Modal using the following steps:
78 | 
79 | 1. Modify you `.env.production` file in `frontend` dir and add your `VITE_BACKEND_URL` for the production environment
80 | 2. Build the frontend `npm run build` - this will create a `dist` folder with the frontend bundle
81 | 3. `modal deploy .\main.py`
82 |

--------------------------------------------------------------------------------
/main.py:
--------------------------------------------------------------------------------
  1 | import json
  2 | from pathlib import Path
  3 | from modal import Secret, asgi_app, Image, concurrent
  4 | from vrag.app import app
  5 | from vrag.colpali import ColPaliModel
  6 | 
  7 | static_path = Path(__file__).with_name("frontend").joinpath("dist").resolve()
  8 | 
  9 | img = (
 10 |     Image.debian_slim(python_version="3.11")
 11 |     .pip_install(
 12 |         "openai==1.44.1",
 13 |         "opencv_python_headless==4.10.0.84",
 14 |         "pydantic==2.9.1",
 15 |         "pypdfium2==4.30.0",
 16 |         "fastapi==0.114.2",
 17 |         "qdrant_client==1.11.2",
 18 |         "sse-starlette==2.1.3",
 19 |     )
 20 |     .pip_install("numpy==2.1.1")
 21 |     .add_local_python_source("vrag")
 22 |     .add_local_dir(static_path, remote_path="/assets")
 23 | )
 24 | 
 25 | colpali = ColPaliModel()
 26 | 
 27 | 
 28 | @app.function(
 29 |     image=img,
 30 |     secrets=[Secret.from_dotenv()],
 31 |     max_containers=1,
 32 |     scaledown_window=300,
 33 |     timeout=600,
 34 | )
 35 | @concurrent(max_inputs=10)
 36 | @asgi_app()
 37 | def web():
 38 |     from uuid import UUID
 39 |     import uuid
 40 |     from pydantic import BaseModel
 41 |     from fastapi import FastAPI, UploadFile, File
 42 |     from sse_starlette.sse import EventSourceResponse, ServerSentEvent
 43 |     from fastapi.staticfiles import StaticFiles
 44 |     from fastapi.middleware.cors import CORSMiddleware
 45 |     from vrag.vrag import VRAG
 46 |     from vrag.qdrant_client import InMemoryQdrant
 47 | 
 48 |     class SearchRequest(BaseModel):
 49 |         query: str
 50 |         instance_id: UUID
 51 |         count: int = 3
 52 | 
 53 |     web_app = FastAPI()
 54 | 
 55 |     origins = [
 56 |         "http://localhost",
 57 |         "http://localhost:5173",
 58 |     ]
 59 | 
 60 |     web_app.add_middleware(
 61 |         CORSMiddleware,
 62 |         allow_origins=origins,
 63 |         allow_credentials=True,
 64 |         allow_methods=["*"],
 65 |         allow_headers=["*"],
 66 |     )
 67 | 
 68 |     qdrant = InMemoryQdrant()
 69 | 
 70 |     VisionRAG = VRAG(colpali=colpali, qdrant=qdrant)
 71 | 
 72 |     @web_app.post("/collections")
 73 |     async def create_collection(files: list[UploadFile] = File(...)):
 74 |         name = str(uuid.uuid4())
 75 |         filenames = []
 76 |         byte_files = []
 77 | 
 78 |         async def read_files():
 79 |             for file in files:
 80 |                 content = await file.read()
 81 |                 filenames.append(file.filename or "file has no name")
 82 |                 byte_files.append((name, file.filename or "file has no name", content))
 83 | 
 84 |         await read_files()
 85 | 
 86 |         async def event_generator():
 87 |             yield ServerSentEvent(
 88 |                 data=json.dumps({"message": f"Indexing {len(byte_files)} files"})
 89 |             )
 90 |             for idx, byte_file in enumerate(byte_files):
 91 |                 yield ServerSentEvent(
 92 |                     data=json.dumps(
 93 |                         {"message": f"Indexing file {idx + 1} / {len(byte_files)}"}
 94 |                     )
 95 |                 )
 96 |                 try:
 97 |                     async for state in VisionRAG.add_pdf(*byte_file):
 98 |                         yield state
 99 |                 except Exception as e:
100 |                     yield json.dumps({"error": str(e)})
101 |             yield ServerSentEvent(
102 |                 data=json.dumps({"id": name, "filenames": filenames}), event="complete"
103 |             )
104 | 
105 |         return EventSourceResponse(event_generator())
106 | 
107 |     @web_app.post("/search")
108 |     async def search(query: SearchRequest):
109 |         can_query = await qdrant.does_collection_exist(str(query.instance_id))
110 | 
111 |         async def event_generator():
112 |             if not can_query:
113 |                 yield ServerSentEvent(
114 |                     data=json.dumps(
115 |                         {
116 |                             "message": "The index has been deleted or does not exist. Please re-add the files."
117 |                         }
118 |                     )
119 |                 )
120 |                 return
121 |             async for stage in VisionRAG.run_vrag(str(query.instance_id), query.query, query.count):
122 |                 yield stage
123 | 
124 |         return EventSourceResponse(event_generator())
125 | 
126 |     web_app.mount("/", StaticFiles(directory="/assets", html=True))
127 |     return web_app
128 |

--------------------------------------------------------------------------------
/vrag/vrag.py:
--------------------------------------------------------------------------------
  1 | import base64
  2 | import json
  3 | import os
  4 | from typing import AsyncGenerator
  5 | 
  6 | import cv2
  7 | from vrag.colpali import ColPaliModel
  8 | from vrag.qdrant_client import InMemoryQdrant, Result
  9 | from vrag.pdf_to_image import images_from_pdf_bytes
 10 | from openai import OpenAI
 11 | from openai.types.chat import ChatCompletionMessageParam
 12 | from sse_starlette.sse import ServerSentEvent
 13 | 
 14 | 
 15 | class VRAG:
 16 |     def __init__(self, colpali: ColPaliModel, qdrant: InMemoryQdrant):
 17 |         self.colpali = colpali
 18 |         self.qdrant = qdrant
 19 |         self.openai = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
 20 | 
 21 |     async def run_vrag(
 22 |         self, collection_name: str, query: str, count: int = 3
 23 |     ) -> AsyncGenerator[ServerSentEvent, None]:
 24 |         results = await self.search(collection_name, query, count)
 25 | 
 26 |         yield ServerSentEvent(
 27 |             data=json.dumps(
 28 |                 {
 29 |                     "results": [
 30 |                         {
 31 |                             "score": result.score,
 32 |                             "image": result.payload["image"],
 33 |                             "page": result.payload["page"],
 34 |                             "name": result.payload["name"],
 35 |                         }
 36 |                         for result in results
 37 |                     ]
 38 |                 }
 39 |             ),
 40 |             event="sources",
 41 |         )
 42 | 
 43 |         augmented = await self.augment(results, query)
 44 | 
 45 |         async for completion in self.generate(augmented):
 46 |             yield completion
 47 | 
 48 |         imgs: list[str] = [result.payload["image"] for result in results]
 49 |         heatmaps = self.colpali.generate_heatmaps.remote(imgs, query)
 50 |         yield ServerSentEvent(
 51 |             data=json.dumps(heatmaps),
 52 |             event="heatmaps",
 53 |         )
 54 | 
 55 |     # index data
 56 |     async def add_pdf(
 57 |         self, collection_name: str, name: str, pdf: bytes
 58 |     ) -> AsyncGenerator[ServerSentEvent, None]:
 59 |         await self.qdrant.create_collection(collection_name)
 60 | 
 61 |         embeddings: list[list[list[float]]] = []
 62 |         idx = 0
 63 |         batch_size = 4
 64 |         images = images_from_pdf_bytes(pdf)
 65 |         count = len(images)
 66 | 
 67 |         yield ServerSentEvent(
 68 |             data=json.dumps({"message": f"0 % of {count} pages indexed...\n"})
 69 |         )
 70 | 
 71 |         async for embedding in self.colpali.embed_images.remote_gen.aio(
 72 |             images, batch_size
 73 |         ):
 74 |             embeddings.append(embedding)
 75 |             if idx < count:
 76 |                 percent = int(idx / count * 100)
 77 |                 yield ServerSentEvent(
 78 |                     data=json.dumps(
 79 |                         {"message": f"{percent} % of {count} pages indexed...\n"}
 80 |                     )
 81 |                 )
 82 |             idx += 1
 83 | 
 84 |         encoded_images = []
 85 | 
 86 |         for img in images:
 87 |             img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
 88 |             img = cv2.imencode(".jpg", img)[1]
 89 |             img = base64.b64encode(img).decode("utf-8")  # type: ignore
 90 |             encoded_images.append(img)
 91 | 
 92 |         await self.qdrant.upsert_points(
 93 |             collection_name, name, embeddings, encoded_images
 94 |         )
 95 | 
 96 |     # retrieve data
 97 |     async def search(
 98 |         self, collection_name: str, query: str, count: int
 99 |     ) -> list[Result]:
100 |         query_vector = await self.colpali.embed_queries.remote.aio([query])
101 |         results = await self.qdrant.search_collection(
102 |             collection_name, query_vector, count
103 |         )
104 |         return results
105 | 
106 |     # augment
107 |     async def augment(
108 |         self,
109 |         results: list[Result],
110 |         query: str,
111 |     ) -> list[ChatCompletionMessageParam]:
112 | 
113 |         images = []
114 |         for result in results:
115 |             if result and result.payload and "image" in result.payload:
116 |                 image = {
117 |                     "type": "image_url",
118 |                     "image_url": {
119 |                         "url": f"data:image/jpeg;base64,{result.payload['image']}"
120 |                     },
121 |                 }
122 |                 images.append(image)
123 | 
124 |         payload = [
125 |             {
126 |                 "role": "system",
127 |                 "content": "Your task is to answer to the user question based on the provided image or images.",
128 |             },
129 |             {
130 |                 "role": "user",
131 |                 "content": [
132 |                     {"type": "text", "text": query},
133 |                 ]
134 |                 + images,
135 |             },
136 |         ]
137 | 
138 |         return payload
139 | 
140 |     # generate
141 |     async def generate(
142 |         self, query: list[ChatCompletionMessageParam]
143 |     ) -> AsyncGenerator[ServerSentEvent, None]:
144 |         stream = self.openai.chat.completions.create(
145 |             model="gpt-4o", messages=query, stream=True
146 |         )
147 | 
148 |         for chunk in stream:
149 |             if chunk.choices[0].delta.content is not None:
150 |                 yield ServerSentEvent(
151 |                     data=json.dumps({"chunk": chunk.choices[0].delta.content})
152 |                 )
153 |

--------------------------------------------------------------------------------
/vrag/colpali.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | from typing import AsyncGenerator, cast
  3 | import modal
  4 | from vrag.app import app
  5 | 
  6 | 
  7 | img = (
  8 |     modal.Image.debian_slim(python_version="3.11")
  9 |     .pip_install(
 10 |         "colpali_engine==0.3.0",
 11 |         "torch",
 12 |         "transformers==4.44.2",
 13 |         "einops==0.8.0",
 14 |         "vidore_benchmark==4.0.1",
 15 |     )
 16 |     .pip_install("numpy==2.1.1")
 17 |     .pip_install("opencv_python_headless==4.10.0.84")
 18 | )
 19 | 
 20 | 
 21 | class NumpyDataset:
 22 |     def __init__(self, numpy_list: list):
 23 |         import numpy as np
 24 |         import torch
 25 | 
 26 |         self.np = np
 27 |         self.torch = torch
 28 |         self.numpy_list = numpy_list
 29 | 
 30 |     def __len__(self) -> int:
 31 |         return len(self.numpy_list)
 32 | 
 33 |     def __getitem__(self, idx):
 34 |         sample = self.numpy_list[idx]
 35 |         sample = self.torch.from_numpy(sample)
 36 |         return sample
 37 | 
 38 | 
 39 | class StringListDataset:
 40 |     def __init__(self, string_list: list[str]):
 41 |         self.string_list = string_list
 42 | 
 43 |     def __len__(self) -> int:
 44 |         return len(self.string_list)
 45 | 
 46 |     def __getitem__(self, idx) -> str:
 47 |         sample = self.string_list[idx]
 48 |         return sample
 49 | 
 50 | 
 51 | def create_numpy_dataset_class(numpy_list: list):
 52 |     from torch.utils.data import Dataset
 53 | 
 54 |     class DynamicNumpyDataset(NumpyDataset, Dataset):
 55 |         pass
 56 | 
 57 |     return DynamicNumpyDataset(numpy_list)
 58 | 
 59 | 
 60 | def create_stringlist_dataset_class(string_list: list[str]):
 61 |     from torch.utils.data import Dataset
 62 | 
 63 |     class DynamicStringListDataset(StringListDataset, Dataset):
 64 |         pass
 65 | 
 66 |     return DynamicStringListDataset(string_list)
 67 | 
 68 | 
 69 | @app.cls(
 70 |     gpu="A10G",
 71 |     secrets=[modal.Secret.from_dotenv()],
 72 |     cpu=4,
 73 |     timeout=600,
 74 |     scaledown_window=300,
 75 |     image=img,
 76 | )
 77 | class ColPaliModel:
 78 |     def __init__(self):
 79 |         from transformers import PreTrainedModel
 80 |         from colpali_engine.models import (
 81 |             ColPaliProcessor,
 82 |         )
 83 | 
 84 |         self.model_name = "vidore/colpali-v1.2"
 85 |         self.model: PreTrainedModel
 86 |         self.token = os.environ.get("HF_TOKEN")
 87 |         self.processor: ColPaliProcessor
 88 |         self.mock_image = self.create_mock_image()
 89 | 
 90 |     @modal.build()
 91 |     @modal.enter()
 92 |     def load_model(self):
 93 |         import torch
 94 |         from colpali_engine.models import ColPali
 95 |         from colpali_engine.models import (
 96 |             ColPaliProcessor,
 97 |         )
 98 | 
 99 |         if torch.cuda.is_available() and torch.cuda.mem_get_info()[1] >= 8 * 1024**3:
100 |             device = torch.device("cuda")
101 |             torch_type = torch.bfloat16
102 |         else:
103 |             device = torch.device("cpu")
104 |             torch_type = torch.float32
105 | 
106 |         self.model = ColPali.from_pretrained(
107 |             self.model_name,
108 |             torch_dtype=torch_type,
109 |             device_map=device,
110 |             token=self.token,
111 |         ).eval()
112 | 
113 |         self.processor = cast(
114 |             ColPaliProcessor,
115 |             ColPaliProcessor.from_pretrained(self.model_name),
116 |         )
117 | 
118 |     def create_mock_image(self):
119 |         import numpy as np
120 | 
121 |         """Creates a blank 448x448 RGB image."""
122 |         return 255 * np.ones((448, 448, 3), dtype=np.uint8)
123 | 
124 |     def process_images(self, images: list):
125 |         """This function is modified version of the function in colpali_processing_utils.py to remove PIL dependency."""
126 |         texts_doc = ["Describe the image."] * len(images)
127 |         batch_doc = self.processor(
128 |             text=texts_doc,
129 |             images=images,
130 |             return_tensors="pt",
131 |             padding="longest",
132 |         )
133 |         return batch_doc
134 | 
135 |     def process_queries(
136 |         self,
137 |         queries: list[str],
138 |         max_length: int = 50,
139 |         suffix: str | None = None,
140 |     ):
141 |         """This function is modified version of the function in colpali_processing_utils.py to remove PIL dependency."""
142 | 
143 |         if suffix is None:
144 |             suffix = "<pad>" * 10
145 |         texts_query: list[str] = []
146 | 
147 |         for query in queries:
148 |             query = f"Question: {query}"
149 |             query += suffix  # add suffix (pad tokens)
150 |             texts_query.append(query)
151 | 
152 |         batch_query = self.processor(
153 |             images=[self.mock_image] * len(texts_query),
154 |             text=texts_query,
155 |             return_tensors="pt",
156 |             padding="longest",
157 |             max_length=max_length + self.processor.image_seq_length,
158 |         )
159 | 
160 |         del batch_query["pixel_values"]
161 | 
162 |         batch_query["input_ids"] = batch_query["input_ids"][
163 |             ..., self.processor.image_seq_length :
164 |         ]
165 |         batch_query["attention_mask"] = batch_query["attention_mask"][
166 |             ..., self.processor.image_seq_length :
167 |         ]
168 | 
169 |         return batch_query
170 | 
171 |     @modal.method()
172 |     def embed_queries(self, queries: list[str]) -> list[list[float]]:
173 |         import torch
174 | 
175 |         from torch.utils.data import DataLoader
176 | 
177 |         dataloader = DataLoader(
178 |             create_stringlist_dataset_class(queries),
179 |             batch_size=1,
180 |             shuffle=False,
181 |             collate_fn=lambda x: self.process_queries(x),
182 |         )
183 | 
184 |         query_embeddings: list[torch.Tensor] = []
185 | 
186 |         for batch_query in dataloader:
187 |             with torch.no_grad():
188 |                 batch_query = {
189 |                     k: v.to(self.model.device) for k, v in batch_query.items()
190 |                 }
191 |                 embedding_query = self.model(**batch_query)
192 |             query_embeddings.extend(torch.unbind(embedding_query.to("cpu")))
193 | 
194 |         embeddings = query_embeddings[0].tolist()
195 |         return embeddings
196 | 
197 |     @modal.method()
198 |     async def embed_images(
199 |         self, images: list, batch_size: int
200 |     ) -> AsyncGenerator[list[list[float]], None]:
201 |         import torch
202 |         from torch.utils.data import DataLoader
203 | 
204 |         dataloader = DataLoader(
205 |             create_numpy_dataset_class(images),
206 |             batch_size=batch_size,
207 |             shuffle=False,
208 |             collate_fn=lambda x: self.process_images(x),
209 |         )
210 | 
211 |         for batch_doc in dataloader:
212 |             with torch.no_grad():
213 |                 batch_doc = {k: v.to(self.model.device) for k, v in batch_doc.items()}
214 |                 embeddings_doc = self.model(**batch_doc)
215 | 
216 |                 # Shape: [batch_size, 1030, 128]
217 |                 # Convert embeddings_doc to CPU, then iterate over the batch dimension
218 |                 embeddings_list: list[list[list[float]]] = embeddings_doc.to(
219 |                     "cpu"
220 |                 ).tolist()
221 | 
222 |             # Yield each image's embeddings as a list of lists (1030 embeddings of 128 dimensions each)
223 |             for embedding_per_image in embeddings_list:
224 |                 yield embedding_per_image
225 | 
226 |     @modal.method()
227 |     def generate_heatmaps(self, images: list[str], query: str):
228 |         import cv2
229 |         import numpy as np
230 |         import base64
231 | 
232 |         heatmaps: list = []
233 | 
234 |         for image in images:
235 |             b = base64.b64decode(image)
236 |             data = np.frombuffer(b, np.uint8)
237 |             img = cv2.imdecode(data, cv2.IMREAD_COLOR)
238 |             rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
239 |             output_text, query_tokens_list = self.process_text_similarity(query)
240 |             attention_json = self.generate_interpretability_json(
241 |                 output_text, query_tokens_list, rgb
242 |             )
243 |             heatmaps.append(attention_json)
244 | 
245 |         return {
246 |             "heatmaps": heatmaps,
247 |             "query_tokens": [token for token in query_tokens_list if token != ""],
248 |         }
249 | 
250 |     @staticmethod
251 |     def is_special_token(token: str) -> bool:
252 |         # Check if the token meets any of the special conditions
253 |         if len(token) < 2:
254 |             return True
255 |         if token.startswith("<"):
256 |             return True
257 |         if token.isdigit():
258 |             return True
259 |         if token.isspace():
260 |             return True
261 |         if token == "Question":
262 |             return True
263 |         return False
264 | 
265 |     def process_text_similarity(self, query):
266 |         import torch
267 | 
268 |         input_text_processed = self.processor.process_queries([query]).to(
269 |             self.model.device
270 |         )
271 | 
272 |         with torch.no_grad():
273 |             output_text = self.model.forward(
274 |                 **input_text_processed
275 |             )  # (1, query_tokens, dim)
276 | 
277 |         query_tokens_list = self.processor.tokenizer.tokenize(  # type: ignore
278 |             self.processor.decode(input_text_processed.input_ids[0])
279 |         )
280 | 
281 |         # Mark special tokens as empty strings
282 |         query_tokens_list = [
283 |             "" if ColPaliModel.is_special_token(token) else token
284 |             for token in query_tokens_list
285 |         ]
286 | 
287 |         return output_text, query_tokens_list
288 | 
289 |     def generate_interpretability_json(
290 |         self, output_text, query_tokens_list, image
291 |     ) -> list:
292 |         """
293 |         Generates attention heatmap JSON for a query and image pair using the model.
294 |         Modified version of the function gen_and_save_similarity_map_per_token from vidore_benchmark repo.
295 |         See: https://github.com/illuin-tech/vidore-benchmark/blob/main/src/vidore_benchmark/interpretability/gen_similarity_maps.py
296 |         """
297 |         import cv2
298 |         from einops import rearrange
299 |         from vidore_benchmark.interpretability.torch_utils import (
300 |             normalize_similarity_map_per_query_token,
301 |         )
302 |         import torch
303 |         import numpy.typing as npt
304 | 
305 |         patch_size = 14
306 |         resolution = 448
307 |         n_patch_per_dim = resolution // patch_size
308 |         input_image_square = cv2.resize(image, (resolution, resolution))
309 |         input_image_processed = self.process_images([input_image_square]).to(
310 |             self.model.device
311 |         )
312 | 
313 |         with torch.no_grad():
314 |             output_image = self.model.forward(
315 |                 **input_image_processed
316 |             )  # (1, n_patches_x * n_patches_y, dim)
317 | 
318 |         # Remove the special tokens from the output
319 |         output_image = output_image[
320 |             :, : self.processor.image_seq_length, :
321 |         ]  # (1, n_patches_x * n_patches_y, dim)
322 | 
323 |         # Rearrange the output image tensor to explicitly represent the 2D grid of patches
324 |         output_image = rearrange(
325 |             output_image, "b (h w) c -> b h w c", h=n_patch_per_dim, w=n_patch_per_dim
326 |         )  # (1, n_patches_x, n_patches_y, dim)
327 | 
328 |         # Get the similarity map
329 |         similarity_map = torch.einsum(
330 |             "bnk,bijk->bnij", output_text, output_image
331 |         )  # (1, query_tokens, n_patches_x, n_patches_y)
332 | 
333 |         # Normalize the similarity map
334 |         similarity_map_normalized = normalize_similarity_map_per_query_token(
335 |             similarity_map
336 |         )  # (1, query_tokens, n_patches_x, n_patches_y)
337 | 
338 |         # Prepare the attention map data as a JSON structure
339 |         attention_data = []
340 |         sim_map = cast(
341 |             npt.NDArray, similarity_map_normalized.to(torch.float32).cpu().numpy()
342 |         )
343 | 
344 |         # Iterate over the tokens and collect similarity map data for each token
345 |         for token_idx, token in enumerate(query_tokens_list):
346 |             # skip special tokens
347 |             if token == "":
348 |                 continue
349 |             token_data = {
350 |                 "token": query_tokens_list[token_idx],
351 |                 "attention_map": sim_map[0, token_idx, :, :].tolist(),
352 |             }
353 |             attention_data.append(token_data)
354 | 
355 |         return attention_data
356 |

--------------------------------------------------------------------------------
/frontend/src/App.tsx:
--------------------------------------------------------------------------------
  1 | import "@mantine/core/styles.css";
  2 | import "@mantine/dropzone/styles.css";
  3 | import "./styles/global.css";
  4 | import {
  5 |   Badge,
  6 |   Box,
  7 |   Button,
  8 |   Card,
  9 |   Divider,
 10 |   Group,
 11 |   Loader,
 12 |   LoadingOverlay,
 13 |   MantineProvider,
 14 |   Image as MantineImage,
 15 |   Stack,
 16 |   Text,
 17 |   ScrollArea,
 18 |   Skeleton,
 19 |   Flex,
 20 |   Slider,
 21 |   Chip,
 22 |   SegmentedControl,
 23 | } from "@mantine/core";
 24 | import { theme } from "./theme";
 25 | import Header from "./Header/Header";
 26 | import { DropzoneBox } from "./Dropzone/Dropzone";
 27 | import { Question } from "./Question/Question";
 28 | import { useEffect, useState } from "react";
 29 | import ReactMarkdown from "react-markdown";
 30 | import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 31 | import { dark } from "react-syntax-highlighter/dist/esm/styles/prism";
 32 | import * as api from "./api";
 33 | import { IconClearAll } from "@tabler/icons-react";
 34 | import * as classes from "./App.css";
 35 | import { AttentionMap } from "./AttentionMap/AttentionMap";
 36 | 
 37 | type UploadResponse = {
 38 |   id: string;
 39 |   filenames: string[];
 40 |   message: string;
 41 | };
 42 | 
 43 | type SearchRequest = {
 44 |   query: string;
 45 |   instance_id: string;
 46 |   count: number;
 47 | };
 48 | 
 49 | type SearchResponse = {
 50 |   chunk: string;
 51 | };
 52 | 
 53 | type Source = {
 54 |   score: number;
 55 |   image: string;
 56 |   page: number;
 57 |   name: string;
 58 | };
 59 | 
 60 | type Results = {
 61 |   results: Source[];
 62 | };
 63 | 
 64 | type AttentionMap = {
 65 |   token: string;
 66 |   attention_map: number[][];
 67 | };
 68 | 
 69 | type Heatmaps = {
 70 |   query_tokens: string[];
 71 |   heatmaps: AttentionMap[][];
 72 | };
 73 | 
 74 | export default function App() {
 75 |   const [uploading, setUploading] = useState(false);
 76 |   const [loading, setLoading] = useState(false);
 77 |   const [uploadedFiles, setUploadedFiles] = useState<string[]>([]);
 78 |   const [uploadStatus, setUploadStatus] = useState<string>("");
 79 |   const [collection, setCollection] = useState<string | null>(null);
 80 |   const [response, setResponse] = useState<string | null>(null);
 81 |   const [sources, setSources] = useState<Source[]>([]);
 82 |   const [sourceCount, setSourceCount] = useState<number>(3);
 83 |   const [heatmaps, setHeatmaps] = useState<Heatmaps | null>(null);
 84 |   const [width, setWidth] = useState(0);
 85 |   const [height, setHeight] = useState(0);
 86 |   const [tokenMaps, setTokenMaps] = useState<number[][][][] | null>(null);
 87 |   const [currentMap, setCurrentMap] = useState<number>(0);
 88 |   const [color, setColor] = useState<"red" | "blue" | "green">("red");
 89 |   const [opacity, setOpacity] = useState<number>(60);
 90 | 
 91 |   const handleQuestionSubmit = async (question: string) => {
 92 |     setResponse(null);
 93 |     setLoading(true);
 94 |     setSources([]);
 95 |     setHeatmaps(null);
 96 |     setTokenMaps(null);
 97 |     setCurrentMap(0);
 98 |     await api.postStream<SearchRequest, SearchResponse | Results | string>(
 99 |       "search",
100 |       { query: question, instance_id: collection ?? "", count: sourceCount },
101 |       (data, event) => {
102 |         if (event === "sources") {
103 |           setSources((data as Results).results);
104 |         } else if (event === "heatmaps") {
105 |           setHeatmaps(JSON.parse(data as string) as Heatmaps);
106 |         } else {
107 |           setResponse((old) => {
108 |             const res = data as SearchResponse;
109 |             return old ? `${old}${res.chunk}` : res.chunk;
110 |           });
111 |         }
112 |       }
113 |     );
114 |     setLoading(false);
115 |   };
116 | 
117 |   const handleFilesUpload = async (files: File[]) => {
118 |     setUploading(true);
119 |     const formData = new FormData();
120 |     files.forEach((file) => formData.append("files", file, file.name));
121 |     await api.postFilesStream<UploadResponse>(
122 |       "collections",
123 |       formData,
124 |       (data, event) => {
125 |         if (data.message && data.message !== "") {
126 |           setUploadStatus(data.message);
127 |         }
128 |         if (event === "complete") {
129 |           setUploadedFiles(data.filenames);
130 |           setCollection(data.id);
131 |           setUploadStatus("");
132 |         }
133 |       }
134 |     );
135 |     setUploading(false);
136 |   };
137 | 
138 |   useEffect(() => {
139 |     if (heatmaps) {
140 |       const newTokenMaps: number[][][][] = [];
141 | 
142 |       const mapHeight = 32;
143 |       const mapWidth = 32;
144 | 
145 |       heatmaps.heatmaps.forEach((imageHeatmaps) => {
146 |         const tokenMapsForImage: number[][][] = [];
147 |         const combinedAttentionMap: number[][] = [];
148 | 
149 |         for (let i = 0; i < mapHeight; i++) {
150 |           combinedAttentionMap.push(new Array(mapWidth).fill(0));
151 |         }
152 | 
153 |         imageHeatmaps.forEach((heatmap) => {
154 |           const tokenMap: number[][] = [];
155 |           for (let i = 0; i < mapHeight; i++) {
156 |             tokenMap.push(new Array(mapWidth).fill(0));
157 |           }
158 | 
159 |           heatmap.attention_map.forEach((row, rowIndex) => {
160 |             row.forEach((value, colIndex) => {
161 |               combinedAttentionMap[rowIndex][colIndex] += value;
162 |               tokenMap[rowIndex][colIndex] = value;
163 |             });
164 |           });
165 | 
166 |           tokenMapsForImage.push(tokenMap);
167 |         });
168 | 
169 |         newTokenMaps.push(tokenMapsForImage);
170 |       });
171 | 
172 |       setTokenMaps(newTokenMaps);
173 |     }
174 |   }, [heatmaps, width, height]);
175 | 
176 |   return (
177 |     <MantineProvider theme={theme}>
178 |       <Stack w="100%" h="100%" p={{ base: "xs", sm: "sm" }}>
179 |         <Header />
180 |         <Group grow w="100%" justify="center" mah={"100%"}>
181 |           <Box pos="relative">
182 |             <LoadingOverlay
183 |               visible={uploading}
184 |               zIndex={1000}
185 |               overlayProps={{ radius: "sm", blur: 2 }}
186 |               loaderProps={{
187 |                 children: (
188 |                   <Stack gap="md" justify="center" align="center">
189 |                     <Loader color="pink" />
190 |                     <Text>{uploadStatus}</Text>
191 |                   </Stack>
192 |                 ),
193 |               }}
194 |             />
195 |             {uploadedFiles.length === 0 && (
196 |               <DropzoneBox
197 |                 uploading={uploading}
198 |                 onSubmit={(files) => handleFilesUpload(files)}
199 |               />
200 |             )}
201 |             <Group w="100%" justify="center" align="center" gap="md">
202 |               {uploadedFiles.map((filename) => (
203 |                 <Badge
204 |                   key={filename}
205 |                   size="xl"
206 |                   variant="gradient"
207 |                   gradient={{ from: "blue", to: "red", deg: 88 }}
208 |                 >
209 |                   {filename}
210 |                 </Badge>
211 |               ))}
212 |               {uploadedFiles.length > 0 && (
213 |                 <Button
214 |                   color="orange"
215 |                   radius="xl"
216 |                   variant="light"
217 |                   rightSection={<IconClearAll />}
218 |                   onClick={() => {
219 |                     setUploadedFiles([]);
220 |                     setCollection(null);
221 |                   }}
222 |                 >
223 |                   Reset
224 |                 </Button>
225 |               )}
226 |             </Group>
227 |           </Box>
228 |         </Group>
229 |         <Divider />
230 |         <Flex
231 |           direction={{ base: "column", sm: "row" }}
232 |           w="100%"
233 |           mah="100%"
234 |           h="100%"
235 |           mih={0}
236 |           justify="space-between"
237 |         >
238 |           <Stack
239 |             h={{ base: "50%", sm: "100%" }}
240 |             mah="100%"
241 |             pr={{ base: "none", sm: "md" }}
242 |             w={{ base: "100%", sm: "48%" }}
243 |             miw={{ base: "100%", sm: "48%" }}
244 |             pb={{ base: "sm", md: "none", xl: "none", lg: "none" }}
245 |           >
246 |             <Question
247 |               onSubmit={handleQuestionSubmit}
248 |               loading={loading}
249 |               disabled={uploading || uploadedFiles.length === 0}
250 |             />
251 |             {response && response !== "" && (
252 |               <ScrollArea h="100%" mah="100%">
253 |                 <ReactMarkdown
254 |                   children={response}
255 |                   components={{
256 |                     code(props) {
257 |                       const { children, className, node, ...rest } = props;
258 |                       const match = /language-(\w+)/.exec(className || "");
259 |                       return match ? (
260 |                         <SyntaxHighlighter
261 |                           {...rest}
262 |                           PreTag="div"
263 |                           children={String(children).replace(/\n$/, "")}
264 |                           language={match[1]}
265 |                           style={dark}
266 |                         />
267 |                       ) : (
268 |                         <code {...rest} className={className}>
269 |                           {children}
270 |                         </code>
271 |                       );
272 |                     },
273 |                   }}
274 |                 />
275 |               </ScrollArea>
276 |             )}
277 |           </Stack>
278 |           <Divider
279 |             display={{ base: "none", sm: "block" }}
280 |             orientation="vertical"
281 |             maw="1px"
282 |           />
283 |           <Divider
284 |             display={{ base: "block", sm: "none" }}
285 |             orientation="horizontal"
286 |             mah="1px"
287 |           />
288 |           <Stack
289 |             h="100%"
290 |             mah="100%"
291 |             w={{ base: "100%", sm: "48%" }}
292 |             miw={{ base: "100%", sm: "48%" }}
293 |             gap="md"
294 |             pl="md"
295 |             pr="md"
296 |             pt={{ base: "xs", md: "none", xl: "none", lg: "none" }}
297 |             className={classes.sources}
298 |           >
299 |             <Stack>
300 |               <Group grow justify="space-between" gap="xl">
301 |                 <Stack>
302 |                   <Text fz={{ base: "md", sm: "xl" }}>
303 |                     Using top {sourceCount} matches
304 |                   </Text>
305 |                   <Slider
306 |                     color="orange"
307 |                     defaultValue={3}
308 |                     size={"sm"}
309 |                     step={1}
310 |                     onChange={(value) => setSourceCount(value)}
311 |                     value={sourceCount}
312 |                     min={3}
313 |                     max={10}
314 |                     pb={"xl"}
315 |                     disabled={loading}
316 |                     marks={[
317 |                       { value: 3, label: "3" },
318 |                       { value: 4, label: "4" },
319 |                       { value: 5, label: "5" },
320 |                       { value: 6, label: "6" },
321 |                       { value: 7, label: "7" },
322 |                       { value: 8, label: "8" },
323 |                       { value: 9, label: "9" },
324 |                       { value: 10, label: "10" },
325 |                     ]}
326 |                   />
327 |                 </Stack>
328 |                 <Stack gap="sm">
329 |                   <Text size="sm" fw={500}>
330 |                     Heatmap color
331 |                   </Text>
332 |                   <SegmentedControl
333 |                     value={color}
334 |                     color={color}
335 |                     size="xs"
336 |                     variant="light"
337 |                     onChange={(val) =>
338 |                       setColor(val as "red" | "blue" | "green")
339 |                     }
340 |                     data={[
341 |                       { label: "Red", value: "red" },
342 |                       { label: "Green", value: "green" },
343 |                       { label: "Blue", value: "blue" },
344 |                     ]}
345 |                   />
346 |                   <Slider
347 |                     color="orange"
348 |                     label="Heatmap opacity"
349 |                     size="xs"
350 |                     value={opacity}
351 |                     onChange={(value) => setOpacity(value)}
352 |                     marks={[
353 |                       { value: 20, label: "20%" },
354 |                       { value: 50, label: "50%" },
355 |                       { value: 80, label: "80%" },
356 |                     ]}
357 |                   />
358 |                 </Stack>
359 |               </Group>
360 |               <Group pt="md">
361 |                 {heatmaps?.query_tokens.map((token, idx) => (
362 |                   <Chip
363 |                     key={idx}
364 |                     color="grape"
365 |                     size="sm"
366 |                     checked={idx == currentMap}
367 |                     radius="xl"
368 |                     onChange={() => setCurrentMap(idx)}
369 |                   >
370 |                     {token}
371 |                   </Chip>
372 |                 ))}
373 |               </Group>
374 |             </Stack>
375 |             {sources.length > 0 && (
376 |               <ScrollArea h="100%" mah="100%">
377 |                 {sources.map((source, idx) => (
378 |                   <Card shadow="sm" padding="lg" radius="md" mb="lg" withBorder>
379 |                     <Group justify="space-between" pb="md">
380 |                       <Text fw={500}>
381 |                         {source.name} - Page {source.page}
382 |                       </Text>
383 |                       <Badge color="pink">Score {source.score}</Badge>
384 |                     </Group>
385 |                     <Card.Section style={{ position: "relative" }}>
386 |                       {/* Original Image */}
387 |                       <MantineImage
388 |                         src={"data:image/jpeg;base64," + source.image}
389 |                         height={"100%"}
390 |                         width={"100%"}
391 |                         alt="Original image"
392 |                         onLoad={(img) => {
393 |                           setWidth(img.currentTarget.width);
394 |                           setHeight(img.currentTarget.height);
395 |                         }}
396 |                       />
397 |                       <AttentionMap
398 |                         tokenMap={tokenMaps?.[idx]?.[currentMap] ?? null}
399 |                         width={width}
400 |                         height={height}
401 |                         color={color}
402 |                         opacity={opacity}
403 |                       />
404 |                     </Card.Section>
405 |                   </Card>
406 |                 ))}
407 |               </ScrollArea>
408 |             )}
409 |             {sources.length === 0 && loading && (
410 |               <Skeleton height={"100%"} width={"100%"} visible={loading} />
411 |             )}
412 |           </Stack>
413 |         </Flex>
414 |       </Stack>
415 |     </MantineProvider>
416 |   );
417 | }
418 |

--------------------------------------------------------------------------------