Polarisailabs commited on
Commit
ceb401c
·
verified ·
1 Parent(s): aaf8efa

Upload 3 files

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ docs/REGULATION[[:space:]](EU)[[:space:]]2024:1689[[:space:]]OF[[:space:]]THE[[:space:]]EUROPEAN[[:space:]]PARLIAMENT[[:space:]]AND[[:space:]]OF[[:space:]]THE[[:space:]]COUNCIL.pdf filter=lfs diff=lfs merge=lfs -text
app.py ADDED
@@ -0,0 +1,428 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import io
3
+ import json
4
+ import pathlib
5
+ import shutil
6
+ from typing import List, Tuple, Dict
7
+ import gradio as gr
8
+ import numpy as np
9
+ import faiss
10
+ from sentence_transformers import SentenceTransformer
11
+ from pypdf import PdfReader
12
+ import fitz # PyMuPDF
13
+ from collections import defaultdict
14
+ from openai import OpenAI
15
+
16
+ # =========================
17
+ # LLM Endpoint
18
+ # =========================
19
+
20
+ API_KEY = os.environ.get("API_KEY")
21
+ if not API_KEY:
22
+ raise RuntimeError("Missing API_KEY (set it in Hugging Face: Settings → Variables and secrets).")
23
+ client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=API_KEY)
24
+
25
+ # Friendly labels for dropdown
26
+ MODEL_LABELS = {
27
+ "GPT": "gpt-oss:20b",
28
+ "Deepseek": "deepseek-r1",
29
+ "Gemma": "gemma3:27b",
30
+ "Qwen": "qwen3-235b"
31
+ }
32
+ MODEL_MAPPING = {
33
+ "gpt-oss:20b": "openai/gpt-oss-20b:free",
34
+ "deepseek-r1": "deepseek/deepseek-r1:free",
35
+ "gemma3:27b": "google/gemma-3-27b-it:free",
36
+ "qwen3-235b": "qwen/qwen3-235b-a22b:free"
37
+ }
38
+ DEFAULT_MODEL_LABEL = "Deepseek"
39
+ GEN_TEMPERATURE = 0.2
40
+ GEN_TOP_P = 0.95
41
+ GEN_MAX_TOKENS = 1024
42
+
43
+ EMB_MODEL_NAME = "intfloat/multilingual-e5-base"
44
+
45
+ def choose_store_dir() -> Tuple[str, bool]:
46
+ data_root = "/data"
47
+ if os.path.isdir(data_root) and os.access(data_root, os.W_OK):
48
+ d = os.path.join(data_root, "rag_store")
49
+ try:
50
+ os.makedirs(d, exist_ok=True)
51
+ testf = os.path.join(d, ".write_test")
52
+ with open(testf, "w", encoding="utf-8") as f:
53
+ f.write("ok")
54
+ os.remove(testf)
55
+ return d, True
56
+ except Exception:
57
+ pass
58
+ d = os.path.join(os.getcwd(), "store")
59
+ os.makedirs(d, exist_ok=True)
60
+ return d, False
61
+
62
+ STORE_DIR, IS_PERSISTENT = choose_store_dir()
63
+ META_PATH = os.path.join(STORE_DIR, "meta.json")
64
+ INDEX_PATH = os.path.join(STORE_DIR, "faiss.index")
65
+
66
+ LEGACY_STORE_DIR = os.path.join(os.getcwd(), "store")
67
+
68
+ def migrate_legacy_if_any():
69
+ try:
70
+ if IS_PERSISTENT:
71
+ legacy_meta = os.path.join(LEGACY_STORE_DIR, "meta.json")
72
+ legacy_index = os.path.join(LEGACY_STORE_DIR, "faiss.index")
73
+ if (not os.path.exists(META_PATH) or not os.path.exists(INDEX_PATH)) \
74
+ and os.path.isdir(LEGACY_STORE_DIR) \
75
+ and os.path.exists(legacy_meta) and os.path.exists(legacy_index):
76
+ shutil.copyfile(legacy_meta, META_PATH)
77
+ shutil.copyfile(legacy_index, INDEX_PATH)
78
+ except Exception:
79
+ pass
80
+
81
+ migrate_legacy_if_any()
82
+
83
+ _emb_model = None
84
+ _index: faiss.Index = None
85
+ _meta: Dict[str, Dict] = {}
86
+
87
+ DEFAULT_TOP_K = 6
88
+ DEFAULT_POOL_K = 40
89
+ DEFAULT_PER_SOURCE_CAP = 2
90
+ DEFAULT_STRATEGY = "mmr"
91
+ DEFAULT_MMR_LAMBDA = 0.5
92
+
93
+ def get_emb_model():
94
+ global _emb_model
95
+ if _emb_model is None:
96
+ _emb_model = SentenceTransformer(EMB_MODEL_NAME)
97
+ return _emb_model
98
+
99
+ def _ensure_index(dim: int):
100
+ global _index
101
+ if _index is None:
102
+ _index = faiss.IndexFlatIP(dim)
103
+
104
+ def _persist():
105
+ faiss.write_index(_index, INDEX_PATH)
106
+ with open(META_PATH, "w", encoding="utf-8") as f:
107
+ json.dump(_meta, f, ensure_ascii=False)
108
+
109
+ def _load_if_any():
110
+ global _index, _meta
111
+ if os.path.exists(INDEX_PATH) and os.path.exists(META_PATH):
112
+ _index = faiss.read_index(INDEX_PATH)
113
+ with open(META_PATH, "r", encoding="utf-8") as f:
114
+ _meta = json.load(f)
115
+
116
+ def _chunk_text(text: str, chunk_size: int = 800, overlap: int = 120) -> List[str]:
117
+ text = text.replace("\u0000", "")
118
+ res, i, n = [], 0, len(text)
119
+ while i < n:
120
+ j = min(i + chunk_size, n)
121
+ seg = text[i:j].strip()
122
+ if seg:
123
+ res.append(seg)
124
+ i = max(0, j - overlap)
125
+ if j >= n:
126
+ break
127
+ return res
128
+
129
+ def _read_bytes(file) -> bytes:
130
+ if isinstance(file, dict):
131
+ p = file.get("path") or file.get("name")
132
+ if p and os.path.exists(p):
133
+ with open(p, "rb") as f:
134
+ return f.read()
135
+ if "data" in file and isinstance(file["data"], (bytes, bytearray)):
136
+ return bytes(file["data"])
137
+ if isinstance(file, (str, pathlib.Path)):
138
+ with open(file, "rb") as f:
139
+ return f.read()
140
+ if hasattr(file, "read"):
141
+ try:
142
+ if hasattr(file, "seek"):
143
+ try:
144
+ file.seek(0)
145
+ except Exception:
146
+ pass
147
+ return file.read()
148
+ finally:
149
+ try:
150
+ file.close()
151
+ except Exception:
152
+ pass
153
+ raise ValueError("Unsupported file type from gr.File")
154
+
155
+ def _decode_best_effort(raw: bytes) -> str:
156
+ for enc in ["utf-8", "cp932", "shift_jis", "cp950", "big5", "gb18030", "latin-1"]:
157
+ try:
158
+ return raw.decode(enc)
159
+ except Exception:
160
+ continue
161
+ return raw.decode("utf-8", errors="ignore")
162
+
163
+ def _read_pdf(file_bytes: bytes) -> str:
164
+ try:
165
+ with fitz.open(stream=file_bytes, filetype="pdf") as doc:
166
+ if doc.is_encrypted:
167
+ try:
168
+ doc.authenticate("")
169
+ except Exception:
170
+ pass
171
+ texts = [(page.get_text("text") or "") for page in doc]
172
+ txt = "\n".join(texts)
173
+ if txt.strip():
174
+ return txt
175
+ except Exception:
176
+ pass
177
+ try:
178
+ reader = PdfReader(io.BytesIO(file_bytes))
179
+ pages = []
180
+ for p in reader.pages:
181
+ try:
182
+ pages.append(p.extract_text() or "")
183
+ except Exception:
184
+ pages.append("")
185
+ return "\n".join(pages)
186
+ except Exception:
187
+ return ""
188
+
189
+ def _read_any(file) -> str:
190
+ if isinstance(file, dict):
191
+ name = (file.get("orig_name") or file.get("name") or file.get("path") or "upload").lower()
192
+ else:
193
+ name = getattr(file, "name", None) or (str(file) if isinstance(file, (str, pathlib.Path)) else "upload")
194
+ name = name.lower()
195
+ raw = _read_bytes(file)
196
+ if name.endswith(".pdf"):
197
+ return _read_pdf(raw).replace("\u0000", "")
198
+ return _decode_best_effort(raw).replace("\u0000", "")
199
+
200
+ DOCS_DIR = os.path.join(os.getcwd(), "docs")
201
+
202
+ def get_docs_files() -> List[str]:
203
+ if not os.path.isdir(DOCS_DIR):
204
+ return []
205
+ files = []
206
+ for fname in os.listdir(DOCS_DIR):
207
+ if fname.lower().endswith((".pdf", ".txt")):
208
+ files.append(os.path.join(DOCS_DIR, fname))
209
+ return files
210
+
211
+ def build_corpus_from_docs():
212
+ global _index, _meta
213
+ files = get_docs_files()
214
+ if not files:
215
+ return "No files found in docs folder."
216
+ emb_model = get_emb_model()
217
+ chunks, sources, failed = [], [], []
218
+ _index = None
219
+ _meta = {}
220
+ for f in files:
221
+ fname = os.path.basename(f)
222
+ try:
223
+ text = _read_any(f) or ""
224
+ parts = _chunk_text(text)
225
+ if not parts:
226
+ failed.append(fname)
227
+ continue
228
+ chunks.extend(parts)
229
+ sources.extend([fname] * len(parts))
230
+ except Exception:
231
+ failed.append(fname)
232
+ if not chunks:
233
+ return "No text extracted from docs."
234
+ passages = [f"passage: {c}" for c in chunks]
235
+ vec = emb_model.encode(passages, batch_size=64, convert_to_numpy=True, normalize_embeddings=True)
236
+ _ensure_index(vec.shape[1])
237
+ _index.add(vec)
238
+ for i, (src, c) in enumerate(zip(sources, chunks)):
239
+ _meta[str(i)] = {"source": src, "text": c}
240
+ _persist()
241
+ msg = f"Indexed {len(chunks)} chunks from {len(files)} files."
242
+ if failed:
243
+ msg += f" Failed files: {', '.join(failed)}"
244
+ return msg
245
+
246
+ def _encode_query_vec(query: str) -> np.ndarray:
247
+ return get_emb_model().encode([f"query: {query}"], convert_to_numpy=True, normalize_embeddings=True)
248
+
249
+ def retrieve_candidates(qvec: np.ndarray, pool_k: int = 40) -> List[Tuple[str, float]]:
250
+ if _index is None or _index.ntotal == 0:
251
+ return []
252
+ pool_k = min(pool_k, _index.ntotal)
253
+ D, I = _index.search(qvec, pool_k)
254
+ return [(str(idx), float(score)) for idx, score in zip(I[0], D[0]) if idx != -1]
255
+
256
+ def select_diverse_by_source(cands: List[Tuple[str, float]], top_k: int = 6, per_source_cap: int = 2) -> List[Tuple[str, float]]:
257
+ if not cands:
258
+ return []
259
+ by_src: Dict[str, List[Tuple[str, float]]] = defaultdict(list)
260
+ for cid, s in cands:
261
+ m = _meta.get(cid)
262
+ if not m:
263
+ continue
264
+ by_src[m["source"]].append((cid, s))
265
+ for src in by_src:
266
+ by_src[src] = by_src[src][:per_source_cap]
267
+ picked, src_items, ptrs = [], [(s, it) for s, it in by_src.items()], {s: 0 for s in by_src}
268
+ while len(picked) < top_k:
269
+ advanced = False
270
+ for src, items in src_items:
271
+ i = ptrs[src]
272
+ if i < len(items):
273
+ picked.append(items[i])
274
+ ptrs[src] = i + 1
275
+ advanced = True
276
+ if len(picked) >= top_k:
277
+ break
278
+ if not advanced:
279
+ break
280
+ if len(picked) < top_k:
281
+ seen = {cid for cid, _ in picked}
282
+ for cid, s in cands:
283
+ if cid not in seen:
284
+ picked.append((cid, s))
285
+ seen.add(cid)
286
+ if len(picked) >= top_k:
287
+ break
288
+ return picked[:top_k]
289
+
290
+ def _encode_chunks_text(cids: List[str]) -> np.ndarray:
291
+ texts = [f"passage: {(_meta.get(cid) or {}).get('text','')}" for cid in cids]
292
+ return get_emb_model().encode(texts, convert_to_numpy=True, normalize_embeddings=True)
293
+
294
+ def select_diverse_mmr(cands: List[Tuple[str, float]], qvec: np.ndarray, top_k: int = 6, mmr_lambda: float = 0.5) -> List[Tuple[str, float]]:
295
+ if not cands:
296
+ return []
297
+ cids = [cid for cid, _ in cands]
298
+ cvecs = _encode_chunks_text(cids)
299
+ sim_to_q = (cvecs @ qvec.T).reshape(-1)
300
+ selected, remaining = [], set(range(len(cids)))
301
+ while len(selected) < min(top_k, len(cids)):
302
+ if not selected:
303
+ i = int(np.argmax(sim_to_q))
304
+ selected.append(i)
305
+ remaining.remove(i)
306
+ continue
307
+ S = cvecs[selected]
308
+ sim_to_S = (cvecs[list(remaining)] @ S.T)
309
+ max_sim_to_S = sim_to_S.max(axis=1) if sim_to_S.size > 0 else np.zeros((len(remaining),), dtype=np.float32)
310
+ sim_q_rem = sim_to_q[list(remaining)]
311
+ mmr_scores = mmr_lambda * sim_q_rem - (1.0 - mmr_lambda) * max_sim_to_S
312
+ j_rel = int(np.argmax(mmr_scores))
313
+ j = list(remaining)[j_rel]
314
+ selected.append(j)
315
+ remaining.remove(j)
316
+ return [(cids[i], float(sim_to_q[i])) for i in selected][:top_k]
317
+
318
+ def retrieve_diverse(query: str,
319
+ top_k: int = 6,
320
+ pool_k: int = 40,
321
+ per_source_cap: int = 2,
322
+ strategy: str = "mmr",
323
+ mmr_lambda: float = 0.5) -> List[Tuple[str, float]]:
324
+ qvec = _encode_query_vec(query)
325
+ cands = retrieve_candidates(qvec, pool_k=pool_k)
326
+ if strategy == "mmr":
327
+ return select_diverse_mmr(cands, qvec, top_k=top_k, mmr_lambda=mmr_lambda)
328
+ return select_diverse_by_source(cands, top_k=top_k, per_source_cap=per_source_cap)
329
+
330
+ def _format_ctx(hits: List[Tuple[str, float]]) -> str:
331
+ if not hits:
332
+ return ""
333
+ lines = []
334
+ for cid, _ in hits:
335
+ m = _meta.get(cid)
336
+ if not m:
337
+ continue
338
+ source_clean = m.get("source", "")
339
+ text_clean = (m.get("text", "") or "").replace("\n", " ")
340
+ lines.append(f"[{cid}] ({source_clean}) " + text_clean)
341
+ return "\n".join(lines[:10])
342
+
343
+ def chat_fn(message, history, model_label):
344
+ # Map label to model key
345
+ model_key = MODEL_LABELS.get(model_label, MODEL_LABELS[DEFAULT_MODEL_LABEL])
346
+ model_name = MODEL_MAPPING.get(model_key, MODEL_MAPPING[MODEL_LABELS[DEFAULT_MODEL_LABEL]])
347
+ if _index is None or _index.ntotal == 0:
348
+ status = build_corpus_from_docs()
349
+ if not (_index and _index.ntotal > 0):
350
+ yield f"**Index Status:** {status}\n\nPlease ensure you have a 'docs' folder with PDF/TXT files and try again."
351
+ return
352
+ hits = retrieve_diverse(
353
+ message,
354
+ top_k=6,
355
+ pool_k=40,
356
+ per_source_cap=2,
357
+ strategy="mmr",
358
+ mmr_lambda=0.5,
359
+ )
360
+ ctx = _format_ctx(hits) if hits else "(Current index is empty or no matching chunks found)"
361
+ sys_blocks = ["You are a research assistant who has an excellent factual understanding of the policies, regulations, and compliance of enterprises, governments, and global organizations. You are a research assistant who reads policy papers and provides factual answers to queries. If you do not know the answer, you should convey that to the user instead of hallucinating. Answers must be based on retrieved content with evidence and source numbers cited. If retrieval is insufficient, please clearly explain the shortcomings. When answering, please cite the numbers, e.g., [3]"]
362
+ messages = [{"role": "system", "content": "\n\n".join(sys_blocks)}]
363
+ for u, a in history:
364
+ messages.append({"role": "user", "content": u})
365
+ messages.append({"role": "assistant", "content": a})
366
+ messages.append({"role": "user", "content": message})
367
+ try:
368
+ response = client.chat.completions.create(
369
+ model=model_name,
370
+ messages=messages,
371
+ temperature=GEN_TEMPERATURE,
372
+ top_p=GEN_TOP_P,
373
+ max_tokens=GEN_MAX_TOKENS,
374
+ stream=True,
375
+ )
376
+ partial_message = ""
377
+ for chunk in response:
378
+ if hasattr(chunk.choices[0], "delta") and chunk.choices[0].delta.content is not None:
379
+ partial_message += chunk.choices[0].delta.content
380
+ yield partial_message
381
+ elif hasattr(chunk.choices[0], "message") and chunk.choices[0].message.content is not None:
382
+ partial_message += chunk.choices[0].message.content
383
+ yield partial_message
384
+ except Exception as e:
385
+ yield f"[Exception] {repr(e)}"
386
+ with gr.Blocks(theme=gr.themes.Default(primary_hue="sky")) as policyprodigy:
387
+ gr.Markdown("")
388
+ with gr.Row():
389
+ model_dropdown = gr.Dropdown(
390
+ choices=list(MODEL_LABELS.keys()),
391
+ value=DEFAULT_MODEL_LABEL,
392
+ label="Select Model",
393
+ interactive=True
394
+ )
395
+ with gr.Row():
396
+ query_box = gr.Textbox(
397
+ label="Try: ",
398
+ placeholder="Enter your question here...",
399
+ scale=8
400
+ )
401
+ send_btn = gr.Button("Send", scale=1)
402
+ with gr.Row():
403
+ chatbot = gr.Chatbot(label="Maris")
404
+
405
+ state = gr.State([])
406
+
407
+ def chat_wrapper(user_message, history, model_label):
408
+ history = history or []
409
+ gen = chat_fn(user_message, history, model_label)
410
+ result = ""
411
+ for chunk in gen:
412
+ result = chunk
413
+ history.append((user_message, result))
414
+ return history, history
415
+
416
+ send_btn.click(
417
+ chat_wrapper,
418
+ inputs=[query_box, state, model_dropdown],
419
+ outputs=[chatbot, state]
420
+ )
421
+
422
+ try:
423
+ _load_if_any()
424
+ except Exception:
425
+ pass
426
+
427
+ if __name__ == "__main__":
428
+ policyprodigy.launch()
docs/.DS_Store ADDED
Binary file (6.15 kB). View file
 
docs/REGULATION (EU) 2024:1689 OF THE EUROPEAN PARLIAMENT AND OF THE COUNCIL.pdf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bba630444b3278e881066774002a1d7824308934f49ccfa203e65be43692f55e
3
+ size 2583319