dinhanhx / dinhanhx/gpu-friendly-rag
🔒 Security + ⚡ Performance: trust_remote_code=True, unbatched embedding, and hardcoded query
- 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