Spaces:
Runtime error
Runtime error
Create app-backup.py
Browse files- app-backup.py +639 -0
app-backup.py
ADDED
|
@@ -0,0 +1,639 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import gradio as gr
|
| 2 |
+
import spaces
|
| 3 |
+
import os
|
| 4 |
+
from typing import List, Dict, Any, Optional, Tuple
|
| 5 |
+
import hashlib
|
| 6 |
+
from datetime import datetime
|
| 7 |
+
import numpy as np
|
| 8 |
+
from transformers import pipeline, TextIteratorStreamer
|
| 9 |
+
import torch
|
| 10 |
+
from threading import Thread
|
| 11 |
+
import re
|
| 12 |
+
|
| 13 |
+
# PDF μ²λ¦¬ λΌμ΄λΈλ¬λ¦¬
|
| 14 |
+
try:
|
| 15 |
+
import fitz # PyMuPDF
|
| 16 |
+
PDF_AVAILABLE = True
|
| 17 |
+
except ImportError:
|
| 18 |
+
PDF_AVAILABLE = False
|
| 19 |
+
print("β οΈ PyMuPDF not installed. Install with: pip install pymupdf")
|
| 20 |
+
|
| 21 |
+
try:
|
| 22 |
+
from sentence_transformers import SentenceTransformer
|
| 23 |
+
ST_AVAILABLE = True
|
| 24 |
+
except ImportError:
|
| 25 |
+
ST_AVAILABLE = False
|
| 26 |
+
print("β οΈ Sentence Transformers not installed. Install with: pip install sentence-transformers")
|
| 27 |
+
|
| 28 |
+
# Custom CSS
|
| 29 |
+
custom_css = """
|
| 30 |
+
.gradio-container {
|
| 31 |
+
background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
|
| 32 |
+
min-height: 100vh;
|
| 33 |
+
font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
|
| 34 |
+
}
|
| 35 |
+
|
| 36 |
+
.main-container {
|
| 37 |
+
background: rgba(255, 255, 255, 0.98);
|
| 38 |
+
border-radius: 16px;
|
| 39 |
+
padding: 24px;
|
| 40 |
+
box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
|
| 41 |
+
border: 1px solid rgba(0, 0, 0, 0.05);
|
| 42 |
+
margin: 12px;
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
.pdf-status {
|
| 46 |
+
padding: 12px 16px;
|
| 47 |
+
border-radius: 12px;
|
| 48 |
+
margin: 12px 0;
|
| 49 |
+
font-size: 0.95rem;
|
| 50 |
+
font-weight: 500;
|
| 51 |
+
}
|
| 52 |
+
|
| 53 |
+
.pdf-success {
|
| 54 |
+
background: linear-gradient(135deg, #d4edda 0%, #c3e6cb 100%);
|
| 55 |
+
border: 1px solid #b1dfbb;
|
| 56 |
+
color: #155724;
|
| 57 |
+
}
|
| 58 |
+
|
| 59 |
+
.pdf-error {
|
| 60 |
+
background: linear-gradient(135deg, #f8d7da 0%, #f5c6cb 100%);
|
| 61 |
+
border: 1px solid #f1aeb5;
|
| 62 |
+
color: #721c24;
|
| 63 |
+
}
|
| 64 |
+
|
| 65 |
+
.pdf-info {
|
| 66 |
+
background: linear-gradient(135deg, #d1ecf1 0%, #bee5eb 100%);
|
| 67 |
+
border: 1px solid #9ec5d8;
|
| 68 |
+
color: #0c5460;
|
| 69 |
+
}
|
| 70 |
+
|
| 71 |
+
.rag-context {
|
| 72 |
+
background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
|
| 73 |
+
border-left: 4px solid #f59e0b;
|
| 74 |
+
padding: 12px;
|
| 75 |
+
margin: 12px 0;
|
| 76 |
+
border-radius: 8px;
|
| 77 |
+
font-size: 0.9rem;
|
| 78 |
+
}
|
| 79 |
+
|
| 80 |
+
.thinking-section {
|
| 81 |
+
background: rgba(0, 0, 0, 0.02);
|
| 82 |
+
border: 1px solid rgba(0, 0, 0, 0.1);
|
| 83 |
+
border-radius: 8px;
|
| 84 |
+
padding: 12px;
|
| 85 |
+
margin: 8px 0;
|
| 86 |
+
}
|
| 87 |
+
"""
|
| 88 |
+
|
| 89 |
+
class SimpleTextSplitter:
|
| 90 |
+
"""ν
μ€νΈ λΆν κΈ°"""
|
| 91 |
+
def __init__(self, chunk_size=800, chunk_overlap=100):
|
| 92 |
+
self.chunk_size = chunk_size
|
| 93 |
+
self.chunk_overlap = chunk_overlap
|
| 94 |
+
|
| 95 |
+
def split_text(self, text: str) -> List[str]:
|
| 96 |
+
"""ν
μ€νΈλ₯Ό μ²ν¬λ‘ λΆν """
|
| 97 |
+
chunks = []
|
| 98 |
+
sentences = text.split('. ')
|
| 99 |
+
current_chunk = ""
|
| 100 |
+
|
| 101 |
+
for sentence in sentences:
|
| 102 |
+
if len(current_chunk) + len(sentence) < self.chunk_size:
|
| 103 |
+
current_chunk += sentence + ". "
|
| 104 |
+
else:
|
| 105 |
+
if current_chunk:
|
| 106 |
+
chunks.append(current_chunk.strip())
|
| 107 |
+
current_chunk = sentence + ". "
|
| 108 |
+
|
| 109 |
+
if current_chunk:
|
| 110 |
+
chunks.append(current_chunk.strip())
|
| 111 |
+
|
| 112 |
+
return chunks
|
| 113 |
+
|
| 114 |
+
class PDFRAGSystem:
|
| 115 |
+
"""PDF κΈ°λ° RAG μμ€ν
"""
|
| 116 |
+
|
| 117 |
+
def __init__(self):
|
| 118 |
+
self.documents = {}
|
| 119 |
+
self.document_chunks = {}
|
| 120 |
+
self.embeddings_store = {}
|
| 121 |
+
self.text_splitter = SimpleTextSplitter(chunk_size=800, chunk_overlap=100)
|
| 122 |
+
|
| 123 |
+
# μλ² λ© λͺ¨λΈ μ΄κΈ°ν
|
| 124 |
+
self.embedder = None
|
| 125 |
+
if ST_AVAILABLE:
|
| 126 |
+
try:
|
| 127 |
+
self.embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
|
| 128 |
+
print("β
μλ² λ© λͺ¨λΈ λ‘λ μ±κ³΅")
|
| 129 |
+
except Exception as e:
|
| 130 |
+
print(f"β οΈ μλ² λ© λͺ¨λΈ λ‘λ μ€ν¨: {e}")
|
| 131 |
+
|
| 132 |
+
def extract_text_from_pdf(self, pdf_path: str) -> Dict[str, Any]:
|
| 133 |
+
"""PDFμμ ν
μ€νΈ μΆμΆ"""
|
| 134 |
+
if not PDF_AVAILABLE:
|
| 135 |
+
return {
|
| 136 |
+
"metadata": {
|
| 137 |
+
"title": "PDF Reader Not Available",
|
| 138 |
+
"file_name": os.path.basename(pdf_path),
|
| 139 |
+
"pages": 0
|
| 140 |
+
},
|
| 141 |
+
"full_text": "PDF μ²λ¦¬λ₯Ό μν΄ 'pip install pymupdf'λ₯Ό μ€νν΄μ£ΌμΈμ."
|
| 142 |
+
}
|
| 143 |
+
|
| 144 |
+
try:
|
| 145 |
+
doc = fitz.open(pdf_path)
|
| 146 |
+
text_content = []
|
| 147 |
+
metadata = {
|
| 148 |
+
"title": doc.metadata.get("title", os.path.basename(pdf_path)),
|
| 149 |
+
"pages": len(doc),
|
| 150 |
+
"file_name": os.path.basename(pdf_path)
|
| 151 |
+
}
|
| 152 |
+
|
| 153 |
+
for page_num, page in enumerate(doc):
|
| 154 |
+
text = page.get_text()
|
| 155 |
+
if text.strip():
|
| 156 |
+
text_content.append(text)
|
| 157 |
+
|
| 158 |
+
doc.close()
|
| 159 |
+
|
| 160 |
+
return {
|
| 161 |
+
"metadata": metadata,
|
| 162 |
+
"full_text": "\n\n".join(text_content)
|
| 163 |
+
}
|
| 164 |
+
except Exception as e:
|
| 165 |
+
raise Exception(f"PDF μ²λ¦¬ μ€λ₯: {str(e)}")
|
| 166 |
+
|
| 167 |
+
def process_and_store_pdf(self, pdf_path: str, doc_id: str) -> Dict[str, Any]:
|
| 168 |
+
"""PDF μ²λ¦¬ λ° μ μ₯"""
|
| 169 |
+
try:
|
| 170 |
+
# PDF ν
μ€νΈ μΆμΆ
|
| 171 |
+
pdf_data = self.extract_text_from_pdf(pdf_path)
|
| 172 |
+
|
| 173 |
+
# ν
μ€νΈλ₯Ό μ²ν¬λ‘ λΆν
|
| 174 |
+
chunks = self.text_splitter.split_text(pdf_data["full_text"])
|
| 175 |
+
|
| 176 |
+
if not chunks:
|
| 177 |
+
print("Warning: No chunks created from PDF")
|
| 178 |
+
return {"success": False, "error": "No text content found in PDF"}
|
| 179 |
+
|
| 180 |
+
print(f"Created {len(chunks)} chunks from PDF")
|
| 181 |
+
|
| 182 |
+
# μ²ν¬ μ μ₯
|
| 183 |
+
self.document_chunks[doc_id] = chunks
|
| 184 |
+
|
| 185 |
+
# μλ² λ© μμ± (μ νμ )
|
| 186 |
+
if self.embedder:
|
| 187 |
+
try:
|
| 188 |
+
print("Generating embeddings...")
|
| 189 |
+
embeddings = self.embedder.encode(chunks)
|
| 190 |
+
self.embeddings_store[doc_id] = embeddings
|
| 191 |
+
print(f"Generated {len(embeddings)} embeddings")
|
| 192 |
+
except Exception as e:
|
| 193 |
+
print(f"Warning: Failed to generate embeddings: {e}")
|
| 194 |
+
# μλ² λ© μ€ν¨ν΄λ κ³μ μ§ν
|
| 195 |
+
|
| 196 |
+
# λ¬Έμ μ 보 μ μ₯
|
| 197 |
+
self.documents[doc_id] = {
|
| 198 |
+
"metadata": pdf_data["metadata"],
|
| 199 |
+
"chunk_count": len(chunks),
|
| 200 |
+
"upload_time": datetime.now().isoformat()
|
| 201 |
+
}
|
| 202 |
+
|
| 203 |
+
# λλ²κ·Έ: 첫 λ²μ§Έ μ²ν¬ μΆλ ₯
|
| 204 |
+
print(f"First chunk preview: {chunks[0][:200]}...")
|
| 205 |
+
|
| 206 |
+
return {
|
| 207 |
+
"success": True,
|
| 208 |
+
"doc_id": doc_id,
|
| 209 |
+
"chunks": len(chunks),
|
| 210 |
+
"pages": pdf_data["metadata"]["pages"],
|
| 211 |
+
"title": pdf_data["metadata"]["title"]
|
| 212 |
+
}
|
| 213 |
+
|
| 214 |
+
except Exception as e:
|
| 215 |
+
print(f"Error processing PDF: {e}")
|
| 216 |
+
return {"success": False, "error": str(e)}
|
| 217 |
+
|
| 218 |
+
def search_relevant_chunks(self, query: str, doc_ids: List[str], top_k: int = 3) -> List[Dict]:
|
| 219 |
+
"""κ΄λ ¨ μ²ν¬ κ²μ"""
|
| 220 |
+
all_relevant_chunks = []
|
| 221 |
+
|
| 222 |
+
print(f"Searching chunks for query: '{query[:50]}...' in {len(doc_ids)} documents")
|
| 223 |
+
|
| 224 |
+
# λ¨Όμ λ¬Έμκ° μλμ§ νμΈ
|
| 225 |
+
for doc_id in doc_ids:
|
| 226 |
+
if doc_id not in self.document_chunks:
|
| 227 |
+
print(f"Warning: Document {doc_id} not found in chunks")
|
| 228 |
+
continue
|
| 229 |
+
|
| 230 |
+
chunks = self.document_chunks[doc_id]
|
| 231 |
+
print(f"Document {doc_id} has {len(chunks)} chunks")
|
| 232 |
+
|
| 233 |
+
# μλ² λ© κΈ°λ° κ²μ μλ
|
| 234 |
+
if self.embedder and doc_id in self.embeddings_store:
|
| 235 |
+
try:
|
| 236 |
+
query_embedding = self.embedder.encode([query])[0]
|
| 237 |
+
doc_embeddings = self.embeddings_store[doc_id]
|
| 238 |
+
|
| 239 |
+
# μ½μ¬μΈ μ μ¬λ κ³μ° (μμ νκ²)
|
| 240 |
+
similarities = []
|
| 241 |
+
for i, emb in enumerate(doc_embeddings):
|
| 242 |
+
try:
|
| 243 |
+
query_norm = np.linalg.norm(query_embedding)
|
| 244 |
+
emb_norm = np.linalg.norm(emb)
|
| 245 |
+
|
| 246 |
+
if query_norm > 0 and emb_norm > 0:
|
| 247 |
+
sim = np.dot(query_embedding, emb) / (query_norm * emb_norm)
|
| 248 |
+
similarities.append(sim)
|
| 249 |
+
else:
|
| 250 |
+
similarities.append(0.0)
|
| 251 |
+
except Exception as e:
|
| 252 |
+
print(f"Error calculating similarity for chunk {i}: {e}")
|
| 253 |
+
similarities.append(0.0)
|
| 254 |
+
|
| 255 |
+
# μμ μ²ν¬ μ ν
|
| 256 |
+
if similarities:
|
| 257 |
+
top_indices = np.argsort(similarities)[-min(top_k, len(similarities)):][::-1]
|
| 258 |
+
|
| 259 |
+
for idx in top_indices:
|
| 260 |
+
if idx < len(chunks): # μΈλ±μ€ λ²μ νμΈ
|
| 261 |
+
all_relevant_chunks.append({
|
| 262 |
+
"content": chunks[idx],
|
| 263 |
+
"doc_name": self.documents[doc_id]["metadata"]["file_name"],
|
| 264 |
+
"similarity": similarities[idx]
|
| 265 |
+
})
|
| 266 |
+
print(f"Added chunk {idx} with similarity: {similarities[idx]:.3f}")
|
| 267 |
+
except Exception as e:
|
| 268 |
+
print(f"Error in embedding search: {e}")
|
| 269 |
+
# μλ² λ© μ€ν¨μ ν΄λ°±
|
| 270 |
+
|
| 271 |
+
# μλ² λ©μ΄ μκ±°λ μ€ν¨ν κ²½μ° - κ°λ¨ν μ²μ Nκ° μ²ν¬ λ°ν
|
| 272 |
+
if not all_relevant_chunks:
|
| 273 |
+
print(f"Falling back to simple chunk selection for {doc_id}")
|
| 274 |
+
for i in range(min(top_k, len(chunks))):
|
| 275 |
+
all_relevant_chunks.append({
|
| 276 |
+
"content": chunks[i],
|
| 277 |
+
"doc_name": self.documents[doc_id]["metadata"]["file_name"],
|
| 278 |
+
"similarity": 1.0 - (i * 0.1) # μμλλ‘ κ°μ€μΉ
|
| 279 |
+
})
|
| 280 |
+
print(f"Added chunk {i} (fallback)")
|
| 281 |
+
|
| 282 |
+
# μ μ¬λ κΈ°μ€ μ λ ¬
|
| 283 |
+
all_relevant_chunks.sort(key=lambda x: x.get('similarity', 0), reverse=True)
|
| 284 |
+
|
| 285 |
+
# μμ Kκ° μ ν
|
| 286 |
+
result = all_relevant_chunks[:top_k]
|
| 287 |
+
print(f"Returning {len(result)} chunks")
|
| 288 |
+
|
| 289 |
+
# λλ²κ·Έ: 첫 λ²μ§Έ μ²ν¬ λ΄μ© μΌλΆ μΆλ ₯
|
| 290 |
+
if result:
|
| 291 |
+
print(f"First chunk preview: {result[0]['content'][:100]}...")
|
| 292 |
+
|
| 293 |
+
return result
|
| 294 |
+
|
| 295 |
+
def create_rag_prompt(self, query: str, doc_ids: List[str], top_k: int = 3) -> tuple:
|
| 296 |
+
"""RAG ν둬ννΈ μμ± - 쿼리μ 컨ν
μ€νΈλ₯Ό λΆλ¦¬νμ¬ λ°ν"""
|
| 297 |
+
print(f"Creating RAG prompt for query: '{query[:50]}...' with docs: {doc_ids}")
|
| 298 |
+
|
| 299 |
+
relevant_chunks = self.search_relevant_chunks(query, doc_ids, top_k)
|
| 300 |
+
|
| 301 |
+
if not relevant_chunks:
|
| 302 |
+
print("No relevant chunks found - checking if documents exist")
|
| 303 |
+
# λ¬Έμκ° μλλ° μ²ν¬λ₯Ό λͺ» μ°Ύμ κ²½μ°, 첫 λ²μ§Έ μ²ν¬λΌλ μ¬μ©
|
| 304 |
+
for doc_id in doc_ids:
|
| 305 |
+
if doc_id in self.document_chunks and self.document_chunks[doc_id]:
|
| 306 |
+
print(f"Using first chunk from {doc_id} as fallback")
|
| 307 |
+
relevant_chunks = [{
|
| 308 |
+
"content": self.document_chunks[doc_id][0],
|
| 309 |
+
"doc_name": self.documents[doc_id]["metadata"]["file_name"],
|
| 310 |
+
"similarity": 0.5
|
| 311 |
+
}]
|
| 312 |
+
break
|
| 313 |
+
|
| 314 |
+
if not relevant_chunks:
|
| 315 |
+
print("No documents or chunks available")
|
| 316 |
+
return query, ""
|
| 317 |
+
|
| 318 |
+
print(f"Using {len(relevant_chunks)} chunks for context")
|
| 319 |
+
|
| 320 |
+
# 컨ν
μ€νΈ ꡬμ±
|
| 321 |
+
context_parts = []
|
| 322 |
+
context_parts.append("Based on the following document context, please answer the question below:")
|
| 323 |
+
context_parts.append("=" * 40)
|
| 324 |
+
|
| 325 |
+
for i, chunk in enumerate(relevant_chunks, 1):
|
| 326 |
+
context_parts.append(f"\n[Document Reference {i} - {chunk['doc_name']}]")
|
| 327 |
+
# μ²ν¬ ν¬κΈ° μ¦κ°
|
| 328 |
+
content = chunk['content'][:1000] if len(chunk['content']) > 1000 else chunk['content']
|
| 329 |
+
context_parts.append(content)
|
| 330 |
+
print(f"Added chunk {i} ({len(content)} chars) with similarity: {chunk.get('similarity', 0):.3f}")
|
| 331 |
+
|
| 332 |
+
context_parts.append("\n" + "=" * 40)
|
| 333 |
+
|
| 334 |
+
context = "\n".join(context_parts)
|
| 335 |
+
enhanced_query = f"{context}\n\nQuestion: {query}\n\nAnswer based on the document context provided above:"
|
| 336 |
+
|
| 337 |
+
print(f"Enhanced query length: {len(enhanced_query)} chars (original: {len(query)} chars)")
|
| 338 |
+
|
| 339 |
+
return enhanced_query, context
|
| 340 |
+
|
| 341 |
+
# Initialize model and RAG system
|
| 342 |
+
model_id = "openai/gpt-oss-20b"
|
| 343 |
+
pipe = pipeline(
|
| 344 |
+
"text-generation",
|
| 345 |
+
model=model_id,
|
| 346 |
+
torch_dtype="auto",
|
| 347 |
+
device_map="auto",
|
| 348 |
+
)
|
| 349 |
+
|
| 350 |
+
rag_system = PDFRAGSystem()
|
| 351 |
+
|
| 352 |
+
# Global state for RAG
|
| 353 |
+
rag_enabled = False
|
| 354 |
+
selected_docs = []
|
| 355 |
+
top_k_chunks = 3
|
| 356 |
+
last_context = ""
|
| 357 |
+
|
| 358 |
+
def format_conversation_history(chat_history):
|
| 359 |
+
"""Format conversation history for the model"""
|
| 360 |
+
messages = []
|
| 361 |
+
for item in chat_history:
|
| 362 |
+
role = item["role"]
|
| 363 |
+
content = item["content"]
|
| 364 |
+
if isinstance(content, list):
|
| 365 |
+
content = content[0]["text"] if content and "text" in content[0] else str(content)
|
| 366 |
+
messages.append({"role": role, "content": content})
|
| 367 |
+
return messages
|
| 368 |
+
|
| 369 |
+
@spaces.GPU()
|
| 370 |
+
def generate_response(input_data, chat_history, max_new_tokens, system_prompt, temperature, top_p, top_k, repetition_penalty):
|
| 371 |
+
"""Generate response with optional RAG enhancement"""
|
| 372 |
+
global last_context, rag_enabled, selected_docs, top_k_chunks
|
| 373 |
+
|
| 374 |
+
# Debug logging
|
| 375 |
+
print(f"RAG Enabled: {rag_enabled}")
|
| 376 |
+
print(f"Selected Docs: {selected_docs}")
|
| 377 |
+
print(f"Available Docs: {list(rag_system.documents.keys())}")
|
| 378 |
+
|
| 379 |
+
# Apply RAG if enabled
|
| 380 |
+
if rag_enabled and selected_docs:
|
| 381 |
+
doc_ids = [doc.split(":")[0] for doc in selected_docs]
|
| 382 |
+
enhanced_input, context = rag_system.create_rag_prompt(input_data, doc_ids, top_k_chunks)
|
| 383 |
+
last_context = context
|
| 384 |
+
actual_input = enhanced_input
|
| 385 |
+
print(f"RAG Applied - Original: {len(input_data)} chars, Enhanced: {len(enhanced_input)} chars")
|
| 386 |
+
else:
|
| 387 |
+
actual_input = input_data
|
| 388 |
+
last_context = ""
|
| 389 |
+
print("RAG Not Applied")
|
| 390 |
+
|
| 391 |
+
# Prepare messages
|
| 392 |
+
new_message = {"role": "user", "content": actual_input}
|
| 393 |
+
system_message = [{"role": "system", "content": system_prompt}] if system_prompt else []
|
| 394 |
+
processed_history = format_conversation_history(chat_history)
|
| 395 |
+
messages = system_message + processed_history + [new_message]
|
| 396 |
+
|
| 397 |
+
# Setup streaming
|
| 398 |
+
streamer = TextIteratorStreamer(pipe.tokenizer, skip_prompt=True, skip_special_tokens=True)
|
| 399 |
+
generation_kwargs = {
|
| 400 |
+
"max_new_tokens": max_new_tokens,
|
| 401 |
+
"do_sample": True,
|
| 402 |
+
"temperature": temperature,
|
| 403 |
+
"top_p": top_p,
|
| 404 |
+
"top_k": top_k,
|
| 405 |
+
"repetition_penalty": repetition_penalty,
|
| 406 |
+
"streamer": streamer
|
| 407 |
+
}
|
| 408 |
+
|
| 409 |
+
thread = Thread(target=pipe, args=(messages,), kwargs=generation_kwargs)
|
| 410 |
+
thread.start()
|
| 411 |
+
|
| 412 |
+
# Process streaming output
|
| 413 |
+
thinking = ""
|
| 414 |
+
final = ""
|
| 415 |
+
started_final = False
|
| 416 |
+
|
| 417 |
+
for chunk in streamer:
|
| 418 |
+
if not started_final:
|
| 419 |
+
if "assistantfinal" in chunk.lower():
|
| 420 |
+
split_parts = re.split(r'assistantfinal', chunk, maxsplit=1)
|
| 421 |
+
thinking += split_parts[0]
|
| 422 |
+
final += split_parts[1]
|
| 423 |
+
started_final = True
|
| 424 |
+
else:
|
| 425 |
+
thinking += chunk
|
| 426 |
+
else:
|
| 427 |
+
final += chunk
|
| 428 |
+
|
| 429 |
+
clean_thinking = re.sub(r'^analysis\s*', '', thinking).strip()
|
| 430 |
+
clean_final = final.strip()
|
| 431 |
+
|
| 432 |
+
# Add RAG context indicator if used
|
| 433 |
+
rag_indicator = ""
|
| 434 |
+
if rag_enabled and selected_docs and last_context:
|
| 435 |
+
rag_indicator = "<div class='rag-context'>π RAG Context Applied</div>\n\n"
|
| 436 |
+
|
| 437 |
+
formatted = f"{rag_indicator}<details open><summary>Click to view Thinking Process</summary>\n\n{clean_thinking}\n\n</details>\n\n{clean_final}"
|
| 438 |
+
yield formatted
|
| 439 |
+
|
| 440 |
+
def upload_pdf(file):
|
| 441 |
+
"""PDF νμΌ μ
λ‘λ μ²λ¦¬"""
|
| 442 |
+
if file is None:
|
| 443 |
+
return (
|
| 444 |
+
gr.update(value="<div class='pdf-status pdf-info'>π νμΌμ μ νν΄μ£ΌμΈμ</div>"),
|
| 445 |
+
gr.update(choices=[])
|
| 446 |
+
)
|
| 447 |
+
|
| 448 |
+
try:
|
| 449 |
+
# νμΌ ν΄μλ₯Ό IDλ‘ μ¬μ©
|
| 450 |
+
with open(file.name, 'rb') as f:
|
| 451 |
+
file_hash = hashlib.md5(f.read()).hexdigest()[:8]
|
| 452 |
+
|
| 453 |
+
doc_id = f"doc_{file_hash}"
|
| 454 |
+
|
| 455 |
+
# PDF μ²λ¦¬ λ° μ μ₯
|
| 456 |
+
result = rag_system.process_and_store_pdf(file.name, doc_id)
|
| 457 |
+
|
| 458 |
+
if result["success"]:
|
| 459 |
+
status_html = f"""
|
| 460 |
+
<div class="pdf-status pdf-success">
|
| 461 |
+
β
PDF μ
λ‘λ μλ£!<br>
|
| 462 |
+
π {result['title']}<br>
|
| 463 |
+
π {result['pages']} νμ΄μ§ | π {result['chunks']} μ²ν¬
|
| 464 |
+
</div>
|
| 465 |
+
"""
|
| 466 |
+
|
| 467 |
+
# λ¬Έμ λͺ©λ‘ μ
λ°μ΄νΈ
|
| 468 |
+
doc_choices = [f"{doc_id}: {rag_system.documents[doc_id]['metadata']['file_name']}"
|
| 469 |
+
for doc_id in rag_system.documents.keys()]
|
| 470 |
+
|
| 471 |
+
return (
|
| 472 |
+
status_html,
|
| 473 |
+
gr.update(choices=doc_choices, value=doc_choices)
|
| 474 |
+
)
|
| 475 |
+
else:
|
| 476 |
+
return (
|
| 477 |
+
f"<div class='pdf-status pdf-error'>β μ€λ₯: {result['error']}</div>",
|
| 478 |
+
gr.update()
|
| 479 |
+
)
|
| 480 |
+
|
| 481 |
+
except Exception as e:
|
| 482 |
+
return (
|
| 483 |
+
f"<div class='pdf-status pdf-error'>β μ€λ₯: {str(e)}</div>",
|
| 484 |
+
gr.update()
|
| 485 |
+
)
|
| 486 |
+
|
| 487 |
+
def clear_documents():
|
| 488 |
+
"""λ¬Έμ μ΄κΈ°ν"""
|
| 489 |
+
global selected_docs
|
| 490 |
+
rag_system.documents = {}
|
| 491 |
+
rag_system.document_chunks = {}
|
| 492 |
+
rag_system.embeddings_store = {}
|
| 493 |
+
selected_docs = []
|
| 494 |
+
|
| 495 |
+
return (
|
| 496 |
+
gr.update(value="<div class='pdf-status pdf-info'>ποΈ λͺ¨λ λ¬Έμκ° μμ λμμ΅λλ€</div>"),
|
| 497 |
+
gr.update(choices=[], value=[])
|
| 498 |
+
)
|
| 499 |
+
|
| 500 |
+
def update_rag_settings(enable, docs, k):
|
| 501 |
+
"""Update RAG settings"""
|
| 502 |
+
global rag_enabled, selected_docs, top_k_chunks
|
| 503 |
+
rag_enabled = enable
|
| 504 |
+
selected_docs = docs if docs else []
|
| 505 |
+
top_k_chunks = k
|
| 506 |
+
|
| 507 |
+
# Debug logging
|
| 508 |
+
print(f"RAG Settings Updated - Enabled: {rag_enabled}, Docs: {selected_docs}, Top-K: {top_k_chunks}")
|
| 509 |
+
|
| 510 |
+
status = "β
Enabled" if enable and docs else "β Disabled"
|
| 511 |
+
status_html = f"<div class='pdf-status pdf-info'>π RAG: <strong>{status}</strong></div>"
|
| 512 |
+
|
| 513 |
+
# Show context preview if RAG is enabled
|
| 514 |
+
if enable and docs:
|
| 515 |
+
preview = f"<div class='rag-context'>π Using {len(docs)} document(s) with {k} chunks per query</div>"
|
| 516 |
+
return gr.update(value=status_html), gr.update(value=preview, visible=True)
|
| 517 |
+
else:
|
| 518 |
+
return gr.update(value=status_html), gr.update(value="", visible=False)
|
| 519 |
+
|
| 520 |
+
# Build the interface
|
| 521 |
+
with gr.Blocks(theme=gr.themes.Soft(), css=custom_css, fill_height=True) as demo:
|
| 522 |
+
gr.Markdown("# π GPT-OSS-20B with PDF RAG System")
|
| 523 |
+
gr.Markdown("Enhanced AI assistant with document-based context understanding")
|
| 524 |
+
|
| 525 |
+
with gr.Row():
|
| 526 |
+
# Left sidebar for RAG controls
|
| 527 |
+
with gr.Column(scale=1):
|
| 528 |
+
with gr.Group(elem_classes="main-container"):
|
| 529 |
+
gr.Markdown("### π Document RAG Settings")
|
| 530 |
+
|
| 531 |
+
pdf_upload = gr.File(
|
| 532 |
+
label="Upload PDF",
|
| 533 |
+
file_types=[".pdf"],
|
| 534 |
+
type="filepath"
|
| 535 |
+
)
|
| 536 |
+
|
| 537 |
+
upload_status = gr.HTML(
|
| 538 |
+
value="<div class='pdf-status pdf-info'>π€ Upload a PDF to enable document-based answers</div>"
|
| 539 |
+
)
|
| 540 |
+
|
| 541 |
+
document_list = gr.CheckboxGroup(
|
| 542 |
+
choices=[],
|
| 543 |
+
label="π Select Documents",
|
| 544 |
+
info="Choose documents to use as context"
|
| 545 |
+
)
|
| 546 |
+
|
| 547 |
+
clear_btn = gr.Button("ποΈ Clear All Documents", size="sm", variant="secondary")
|
| 548 |
+
|
| 549 |
+
enable_rag = gr.Checkbox(
|
| 550 |
+
label="β¨ Enable RAG",
|
| 551 |
+
value=False,
|
| 552 |
+
info="Use documents for context-aware responses"
|
| 553 |
+
)
|
| 554 |
+
|
| 555 |
+
top_k_slider = gr.Slider(
|
| 556 |
+
minimum=1,
|
| 557 |
+
maximum=5,
|
| 558 |
+
value=3,
|
| 559 |
+
step=1,
|
| 560 |
+
label="Context Chunks",
|
| 561 |
+
info="Number of document chunks to use"
|
| 562 |
+
)
|
| 563 |
+
|
| 564 |
+
# RAG status display
|
| 565 |
+
rag_status = gr.HTML(
|
| 566 |
+
value="<div class='pdf-status pdf-info'>π RAG: <strong>Disabled</strong></div>"
|
| 567 |
+
)
|
| 568 |
+
|
| 569 |
+
context_preview = gr.HTML(value="", visible=False)
|
| 570 |
+
|
| 571 |
+
# Right side for chat interface
|
| 572 |
+
with gr.Column(scale=3):
|
| 573 |
+
with gr.Group(elem_classes="main-container"):
|
| 574 |
+
# Create ChatInterface with custom function
|
| 575 |
+
chat_interface = gr.ChatInterface(
|
| 576 |
+
fn=generate_response,
|
| 577 |
+
additional_inputs=[
|
| 578 |
+
gr.Slider(label="Max new tokens", minimum=64, maximum=4096, step=1, value=2048),
|
| 579 |
+
gr.Textbox(
|
| 580 |
+
label="System Prompt",
|
| 581 |
+
value="You are a helpful assistant. Reasoning: medium",
|
| 582 |
+
lines=4,
|
| 583 |
+
placeholder="Change system prompt"
|
| 584 |
+
),
|
| 585 |
+
gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=0.7),
|
| 586 |
+
gr.Slider(label="Top-p", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
|
| 587 |
+
gr.Slider(label="Top-k", minimum=1, maximum=100, step=1, value=50),
|
| 588 |
+
gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.0)
|
| 589 |
+
],
|
| 590 |
+
examples=[
|
| 591 |
+
[{"text": "Explain Newton laws clearly and concisely"}],
|
| 592 |
+
[{"text": "Write a Python function to calculate the Fibonacci sequence"}],
|
| 593 |
+
[{"text": "What are the benefits of open weight AI models"}],
|
| 594 |
+
],
|
| 595 |
+
cache_examples=False,
|
| 596 |
+
type="messages",
|
| 597 |
+
description="""Chat with GPT-OSS-20B. Upload PDFs to enhance responses with document context.
|
| 598 |
+
Click to view thinking process (default is on).""",
|
| 599 |
+
textbox=gr.Textbox(
|
| 600 |
+
label="Query Input",
|
| 601 |
+
placeholder="Type your prompt (RAG will be applied if enabled)"
|
| 602 |
+
),
|
| 603 |
+
stop_btn="Stop Generation",
|
| 604 |
+
multimodal=False
|
| 605 |
+
)
|
| 606 |
+
|
| 607 |
+
# Event handlers
|
| 608 |
+
pdf_upload.upload(
|
| 609 |
+
fn=upload_pdf,
|
| 610 |
+
inputs=[pdf_upload],
|
| 611 |
+
outputs=[upload_status, document_list]
|
| 612 |
+
)
|
| 613 |
+
|
| 614 |
+
clear_btn.click(
|
| 615 |
+
fn=clear_documents,
|
| 616 |
+
outputs=[upload_status, document_list]
|
| 617 |
+
)
|
| 618 |
+
|
| 619 |
+
# Update RAG settings when changed
|
| 620 |
+
enable_rag.change(
|
| 621 |
+
fn=update_rag_settings,
|
| 622 |
+
inputs=[enable_rag, document_list, top_k_slider],
|
| 623 |
+
outputs=[rag_status, context_preview]
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
document_list.change(
|
| 627 |
+
fn=update_rag_settings,
|
| 628 |
+
inputs=[enable_rag, document_list, top_k_slider],
|
| 629 |
+
outputs=[rag_status, context_preview]
|
| 630 |
+
)
|
| 631 |
+
|
| 632 |
+
top_k_slider.change(
|
| 633 |
+
fn=update_rag_settings,
|
| 634 |
+
inputs=[enable_rag, document_list, top_k_slider],
|
| 635 |
+
outputs=[rag_status, context_preview]
|
| 636 |
+
)
|
| 637 |
+
|
| 638 |
+
if __name__ == "__main__":
|
| 639 |
+
demo.launch(share=True)
|