dinhanhx / dinhanhx/gpu-friendly-rag

🔒 Security + ⚡ Performance: trust_remote_code=True, unbatched embedding, and hardcoded query

Open
#2 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
0
Forks
0
PR merge metrics
No merged PRs in 30d

Description

## Description

### 1. `trust_remote_code=True` used everywhere — `rag.py` lines 8-12, `setup_db.py` line 10
```python
embedding_model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3", trust_remote_code=True)
rerank_model = AutoModelForSequenceClassification.from_pretrained(..., trust_remote_code=True)
chat_pipe = pipeline(..., trust_remote_code=True)
```
`trust_remote_code=True` executes arbitrary Python from the model repo. This is a known supply-chain risk. For the Jina models it's required, but `Llama-3.2-3B-Instruct` does **not** need it — it uses standard HuggingFace architecture.

### 2. Unbatched embedding in `setup_db.py` lines 18-22
```python
for page in tqdm(pdf):
content = page.get_text()
embedding = embedding_model.encode([content])
embedding_db = np.append(embedding_db, embedding, axis=0)
```
- Each page is encoded individually — no batching. For a GPU pipeline, batching would be significantly faster.
- `np.append` in a loop creates a new array each iteration — O(n²) memory copies. Use a list and `np.vstack` at the end.

### 3. Hardcoded query in `rag.py` line 19
The query is hardcoded. This is fine for a demo/script, but should accept CLI args or stdin for reusability.

### 4. Device handling inconsistency — `rag.py` line 56
```python
device="cuda" if torch.cuda.is_available() else "cpu",
```
But the FAISS index search happens on CPU regardless. The embedding model output may need `.cpu().numpy()` conversion depending on version.

## Suggested Fixes
1. Remove `trust_remote_code=True` from the Llama pipeline
2. Batch PDF pages, encode in groups, collect to list then `np.vstack`
3. Accept query via `sys.argv` or `input()`

Contributor guide

No contributing guide indexed for this repository

Research direction

Read rag.py and setup_db.py at the cited lines, then run the current PDF indexing and query flow to establish its behavior. The work is done when unnecessary remote code trust is removed, PDF embeddings are batched without repeated array appends, queries are accepted as input, and device handling remains compatible with the CPU FAISS search.

Written by the indexing model from the issue text.

Assessment

Tech stack
huggingface, numpy, python, pytorch
Domain
cli, machine-learning, performance, security
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.