damoojeje commited on
Commit
18069c2
·
verified ·
1 Parent(s): e19e3f5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +600 -52
app.py CHANGED
@@ -1,64 +1,612 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
2
- from huggingface_hub import InferenceClient
3
 
4
- """
5
- For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
6
- """
7
- client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 
 
 
 
 
 
8
 
 
 
 
 
9
 
10
- def respond(
11
- message,
12
- history: list[tuple[str, str]],
13
- system_message,
14
- max_tokens,
15
- temperature,
16
- top_p,
17
- ):
18
- messages = [{"role": "system", "content": system_message}]
19
 
20
- for val in history:
21
- if val[0]:
22
- messages.append({"role": "user", "content": val[0]})
23
- if val[1]:
24
- messages.append({"role": "assistant", "content": val[1]})
 
 
 
25
 
26
- messages.append({"role": "user", "content": message})
 
 
 
 
 
 
 
 
 
27
 
28
- response = ""
29
 
30
- for message in client.chat_completion(
31
- messages,
32
- max_tokens=max_tokens,
33
- stream=True,
34
- temperature=temperature,
35
- top_p=top_p,
36
- ):
37
- token = message.choices[0].delta.content
38
 
39
- response += token
40
- yield response
 
 
 
41
 
 
 
 
 
 
 
 
42
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
  """
44
- For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
45
- """
46
- demo = gr.ChatInterface(
47
- respond,
48
- additional_inputs=[
49
- gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
50
- gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
51
- gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
52
- gr.Slider(
53
- minimum=0.1,
54
- maximum=1.0,
55
- value=0.95,
56
- step=0.05,
57
- label="Top-p (nucleus sampling)",
58
- ),
59
- ],
60
- )
61
-
62
-
63
- if __name__ == "__main__":
64
- demo.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import json
3
+ import fitz # PyMuPDF
4
+ import nltk
5
+ import chromadb
6
+ from tqdm import tqdm
7
+ from nltk.tokenize import sent_tokenize
8
+ from sentence_transformers import SentenceTransformer, util
9
+ import numpy as np
10
+ import torch
11
+ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
12
+ import pytesseract
13
+ from PIL import Image
14
+ import io
15
  import gradio as gr
 
16
 
17
+ # ---------------------------
18
+ # ⚙️ Configuration
19
+ # ---------------------------
20
+ pdf_folder = r"./Manuals" # Path relative to the app.py file in the Space
21
+ output_jsonl_pages = "manual_pages_with_ocr.jsonl"
22
+ output_jsonl_chunks = "manual_chunks_with_ocr.jsonl"
23
+ chroma_path = "./chroma_store"
24
+ collection_name = "manual_chunks"
25
+ chunk_size = 750
26
+ chunk_overlap = 100
27
+ MAX_CONTEXT_CHUNKS = 3 # Max chunks to send to the LLM
28
 
29
+ # Hugging Face Model Configuration
30
+ HF_MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
31
+ # Read HF Token from environment variable for security
32
+ HF_TOKEN = os.environ.get("HF_TOKEN") # Hugging Face Space secret name
33
 
34
+ # ---------------------------
35
+ # Ensure NLTK resources are available
36
+ # ---------------------------
37
+ try:
38
+ nltk.data.find('tokenizers/punkt')
39
+ except nltk.downloader.DownloadError:
40
+ nltk.download('punkt')
41
+ except LookupError:
42
+ nltk.download('punkt')
43
 
44
+ # ---------------------------
45
+ # 📄 Utility: Read PDF to text (with OCR fallback)
46
+ # ---------------------------
47
+ # This combines logic from extract_text_from_pdf and extract_text_from_page
48
+ def extract_text_from_page_with_ocr(page):
49
+ text = page.get_text().strip()
50
+ if text:
51
+ return text, False # native text found, no OCR needed
52
 
53
+ # If native text is missing, try OCR
54
+ try:
55
+ pix = page.get_pixmap(dpi=300)
56
+ img_data = pix.tobytes("png")
57
+ img = Image.open(io.BytesIO(img_data))
58
+ ocr_text = pytesseract.image_to_string(img).strip()
59
+ return ocr_text, True
60
+ except Exception as e:
61
+ print(f"OCR failed for a page: {e}")
62
+ return "", False # Return empty and indicate OCR was not used if it fails
63
 
 
64
 
65
+ # ---------------------------
66
+ # 🧹 Clean up lines (from original notebook)
67
+ # ---------------------------
68
+ def clean_text(text):
69
+ lines = text.splitlines()
70
+ lines = [line.strip() for line in lines if line.strip()]
71
+ return "\n".join(lines)
 
72
 
73
+ # ---------------------------
74
+ # ✂️ Sentence Tokenizer (from original notebook)
75
+ # ---------------------------
76
+ def tokenize_sentences(text):
77
+ return sent_tokenize(text)
78
 
79
+ # ---------------------------
80
+ # 📦 Chunk into fixed size blocks (from original notebook)
81
+ # ---------------------------
82
+ def split_into_chunks(sentences, max_tokens=750, overlap=100):
83
+ chunks = []
84
+ current_chunk = []
85
+ current_len = 0
86
 
87
+ for sentence in sentences:
88
+ token_count = len(sentence.split())
89
+ # Check if adding the next sentence exceeds max_tokens
90
+ # If it does, and the current chunk is not empty, save the current chunk
91
+ if current_len + token_count > max_tokens and current_chunk:
92
+ chunks.append(" ".join(current_chunk))
93
+ # Start the next chunk with the overlap
94
+ current_chunk = current_chunk[-overlap:]
95
+ # Recalculate current_len based on the overlap
96
+ current_len = sum(len(s.split()) for s in current_chunk)
97
+
98
+ # Add the current sentence and update length
99
+ current_chunk.append(sentence)
100
+ current_len += token_count
101
+
102
+ # Add the last chunk if it's not empty
103
+ if current_chunk:
104
+ chunks.append(" ".join(current_chunk))
105
+
106
+ return chunks
107
+
108
+
109
+ # ---------------------------
110
+ # 🧠 Extract Metadata from Filename (from original notebook)
111
+ # ---------------------------
112
+ def extract_metadata_from_filename(filename):
113
+ name = filename.lower().replace("_", " ").replace("-", " ")
114
+
115
+ metadata = {
116
+ "model": "unknown",
117
+ "doc_type": "unknown",
118
+ "brand": "life fitness" # Assuming 'life fitness' is constant based on your notebook
119
+ }
120
+
121
+ if "om" in name or "owner" in name:
122
+ metadata["doc_type"] = "owner manual"
123
+ elif "sm" in name or "service" in name:
124
+ metadata["doc_type"] = "service manual"
125
+ elif "assembly" in name:
126
+ metadata["doc_type"] = "assembly instructions"
127
+ elif "alert" in name:
128
+ metadata["doc_type"] = "installer alert"
129
+ elif "parts" in name:
130
+ metadata["doc_type"] = "parts manual"
131
+ elif "bulletin" in name:
132
+ metadata["doc_type"] = "service bulletin"
133
+
134
+ known_models = [
135
+ "se3hd", "se3", "se4", "symbio", "explore", "integrity x", "integrity sl",
136
+ "everest", "engage", "inspire", "discover", "95t", "95x", "95c", "95r", "97c"
137
+ ]
138
+
139
+ for model in known_models:
140
+ # Use regex for more robust matching if needed, but simple 'in' check from notebook
141
+ if model.replace(" ", "") in name.replace(" ", ""):
142
+ metadata["model"] = model
143
+ break
144
+
145
+ return metadata
146
+
147
+ # ---------------------------
148
+ # 🚀 Step 1: Process PDFs, Extract Pages with OCR
149
+ # ---------------------------
150
+ def process_pdfs_for_pages(pdf_folder, output_jsonl):
151
+ print("Starting PDF processing and OCR...")
152
+ all_pages = []
153
+ if not os.path.exists(pdf_folder):
154
+ print(f"Error: PDF folder not found at {pdf_folder}")
155
+ return [] # Return empty list if folder doesn't exist
156
+
157
+ pdf_files = [f for f in os.listdir(pdf_folder) if f.lower().endswith(".pdf")]
158
+ if not pdf_files:
159
+ print(f"No PDF files found in {pdf_folder}")
160
+ return []
161
+
162
+ for pdf_file in tqdm(pdf_files, desc="Scanning PDFs"):
163
+ path = os.path.join(pdf_folder, pdf_file)
164
+ try:
165
+ doc = fitz.open(path)
166
+ for page_num, page in enumerate(doc, start=1):
167
+ text, used_ocr = extract_text_from_page_with_ocr(page)
168
+ if text: # Only save pages with extracted text
169
+ all_pages.append({
170
+ "source_file": pdf_file,
171
+ "page": page_num,
172
+ "text": text,
173
+ "ocr_used": used_ocr
174
+ })
175
+ doc.close() # Close the document
176
+ except Exception as e:
177
+ print(f"Error processing {pdf_file}: {e}")
178
+ continue # Skip to the next file
179
+
180
+ with open(output_jsonl, "w", encoding="utf-8") as f:
181
+ for page in all_pages:
182
+ json.dump(page, f)
183
+ f.write("\n")
184
+
185
+ print(f"✅ Saved {len(all_pages)} pages to {output_jsonl} (with OCR fallback)")
186
+ return all_pages # Return the list of pages
187
+
188
+ # ---------------------------
189
+ # 🚀 Step 2: Chunk the Pages
190
+ # ---------------------------
191
+ def chunk_pages(input_jsonl, output_jsonl, chunk_size, chunk_overlap):
192
+ print("Starting page chunking...")
193
+ all_chunks = []
194
+ if not os.path.exists(input_jsonl):
195
+ print(f"Error: Input JSONL file not found at {input_jsonl}. Run PDF processing first.")
196
+ return []
197
+
198
+ try:
199
+ with open(input_jsonl, "r", encoding="utf-8") as f:
200
+ # Count lines for tqdm progress bar
201
+ total_lines = sum(1 for _ in f)
202
+ f.seek(0) # Reset file pointer to the beginning
203
+
204
+ for line in tqdm(f, total=total_lines, desc="Chunking pages"):
205
+ try:
206
+ page = json.loads(line)
207
+ source_file = page["source_file"]
208
+ page_number = page["page"]
209
+ text = page["text"]
210
+
211
+ metadata = extract_metadata_from_filename(source_file)
212
+ sentences = tokenize_sentences(clean_text(text)) # Clean and tokenize the page text
213
+ chunks = split_into_chunks(sentences, max_tokens=chunk_size, overlap=chunk_overlap)
214
+
215
+ for i, chunk in enumerate(chunks):
216
+ # Ensure chunk text is not empty
217
+ if chunk.strip():
218
+ all_chunks.append({
219
+ "source_file": source_file,
220
+ "chunk_id": f"{source_file}::page_{page_number}::chunk_{i+1}",
221
+ "page": page_number,
222
+ "ocr_used": page.get("ocr_used", False), # Use .get for safety
223
+ "model": metadata.get("model", "unknown"),
224
+ "doc_type": metadata.get("doc_type", "unknown"),
225
+ "brand": metadata.get("brand", "life fitness"),
226
+ "text": chunk.strip() # Ensure no leading/trailing whitespace
227
+ })
228
+ except json.JSONDecodeError:
229
+ print(f"Skipping invalid JSON line: {line}")
230
+ except Exception as e:
231
+ print(f"Error processing page from {line}: {e}")
232
+ continue # Continue with the next line
233
+
234
+ except Exception as e:
235
+ print(f"Error opening or reading input JSONL file: {e}")
236
+ return []
237
+
238
+
239
+ if not all_chunks:
240
+ print("No chunks were created.")
241
+
242
+ with open(output_jsonl, "w", encoding="utf-8") as f:
243
+ for chunk in all_chunks:
244
+ json.dump(chunk, f)
245
+ f.write("\n")
246
+
247
+ print(f"✅ Done! {len(all_chunks)} chunks saved to {output_jsonl}")
248
+ return all_chunks # Return the list of chunks
249
+
250
+ # ---------------------------
251
+ # 🚀 Step 3: Embed Chunks into Chroma
252
+ # ---------------------------
253
+ def embed_chunks_into_chroma(jsonl_path, chroma_path, collection_name):
254
+ print("Starting ChromaDB embedding...")
255
+ try:
256
+ embedder = SentenceTransformer("all-MiniLM-L6-v2")
257
+ embedder.eval()
258
+ print("✅ SentenceTransformer model loaded.")
259
+ except Exception as e:
260
+ print(f"❌ Error loading SentenceTransformer model: {e}")
261
+ return None, "Error loading SentenceTransformer model."
262
+
263
+ try:
264
+ # Use a persistent client
265
+ client = chromadb.PersistentClient(path=chroma_path)
266
+ # Check if collection exists and delete if it does to rebuild
267
+ try:
268
+ client.get_collection(name=collection_name)
269
+ client.delete_collection(collection_name)
270
+ print(f"Deleted existing collection: {collection_name}")
271
+ except Exception: # Collection does not exist, which is fine
272
+ pass
273
+ collection = client.create_collection(name=collection_name)
274
+ print(f"✅ ChromaDB collection '{collection_name}' created.")
275
+ except Exception as e:
276
+ print(f"❌ Error initializing ChromaDB: {e}")
277
+ return None, "Error initializing ChromaDB."
278
+
279
+ texts, metadatas, ids = [], [], []
280
+ batch_size = 16 # Define batch size for embedding
281
+
282
+ if not os.path.exists(jsonl_path):
283
+ print(f"Error: Input JSONL file not found at {jsonl_path}. Run chunking first.")
284
+ return None, "Input chunk file not found."
285
+
286
+ try:
287
+ with open(jsonl_path, "r", encoding="utf-8") as f:
288
+ # Count lines for tqdm progress bar
289
+ total_lines = sum(1 for _ in f)
290
+ f.seek(0) # Reset file pointer to the beginning
291
+
292
+ for line in tqdm(f, total=total_lines, desc="Embedding chunks"):
293
+ try:
294
+ item = json.loads(line)
295
+ texts.append(item.get("text", "")) # Use .get for safety
296
+ ids.append(item.get("chunk_id", f"unknown_{len(ids)}")) # Ensure chunk_id exists
297
+ # Prepare metadata, ensuring all keys are strings and handling potential missing keys
298
+ metadata = {str(k): str(v) for k, v in item.items() if k != "text"}
299
+ metadatas.append(metadata)
300
+
301
+ if len(texts) >= batch_size:
302
+ embeddings = embedder.encode(texts).tolist()
303
+ collection.add(documents=texts, metadatas=metadatas, ids=ids, embeddings=embeddings)
304
+ texts, metadatas, ids = [], [], [] # Reset batches
305
+
306
+ except json.JSONDecodeError:
307
+ print(f"Skipping invalid JSON line during embedding: {line}")
308
+ except Exception as e:
309
+ print(f"Error processing chunk line {line} during embedding: {e}")
310
+ continue # Continue with the next line
311
+
312
+ # Add any remaining items in the last batch
313
+ if texts:
314
+ embeddings = embedder.encode(texts).tolist()
315
+ collection.add(documents=texts, metadatas=metadatas, ids=ids, embeddings=embeddings)
316
+
317
+ print("✅ All OCR-enhanced chunks embedded in Chroma!")
318
+ return collection, None # Return collection and no error
319
+
320
+ except Exception as e:
321
+ print(f"❌ Error reading input JSONL file for embedding: {e}")
322
+ return None, "Error reading input file for embedding."
323
+
324
+
325
+ # ---------------------------
326
+ # 🧠 Load Hugging Face Model and Tokenizer
327
+ # ---------------------------
328
+ # This needs to happen after imports but before the Gradio interface
329
+ tokenizer = None
330
+ model = None
331
+ pipe = None
332
+
333
+ print(f"Attempting to load Hugging Face model: {HF_MODEL_ID}")
334
+ print(f"Using HF_TOKEN (present: {HF_TOKEN is not None})")
335
+
336
+ if not HF_TOKEN:
337
+ print("❌ HF_TOKEN environment variable not set. Cannot load Hugging Face model.")
338
+ else:
339
+ try:
340
+ # Check if CUDA is available
341
+ device = "cuda" if torch.cuda.is_available() else "cpu"
342
+ print(f"Using device: {device}")
343
+
344
+ # Load tokenizer and model
345
+ tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID, token=HF_TOKEN)
346
+ model = AutoModelForCausalLM.from_pretrained(
347
+ HF_MODEL_ID,
348
+ token=HF_TOKEN,
349
+ torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, # Use bfloat16 on GPU
350
+ device_map="auto" if torch.cuda.is_available() else None # Auto device mapping on GPU
351
+ ).to(device) # Move model to selected device
352
+
353
+ # Create a pipeline for easy inference
354
+ pipe = pipeline(
355
+ "text-generation",
356
+ model=model,
357
+ tokenizer=tokenizer,
358
+ max_new_tokens=512,
359
+ temperature=0.1,
360
+ top_p=0.9,
361
+ do_sample=True,
362
+ device=0 if torch.cuda.is_available() else -1 # Specify device for pipeline
363
+ )
364
+
365
+ print(f"✅ Successfully loaded Hugging Face model: {HF_MODEL_ID} on {device}")
366
+
367
+ except Exception as e:
368
+ print(f"❌ Error loading Hugging Face model: {e}")
369
+ print("Please ensure:")
370
+ print("- The HF_TOKEN secret is set in your Hugging Face Space settings.")
371
+ print("- Your Space has sufficient resources (GPU, RAM) for the model.")
372
+ print("- You have accepted the model's terms on Hugging Face (if required).")
373
+ tokenizer, model, pipe = None, None, None # Set to None if loading fails
374
+
375
+
376
+ # ---------------------------
377
+ # 🔎 Query Function (Uses Embedder and Chroma)
378
+ # ---------------------------
379
+ # Embedder is loaded during the embedding step, need to ensure it's accessible
380
+ embedder = None # Initialize embedder as None
381
+
382
+ def query_manuals(question, model_filter=None, doc_type_filter=None, top_k=5, rerank_keywords=None):
383
+ global embedder # Access the global embedder variable
384
+
385
+ if collection is None or embedder is None:
386
+ print("⚠️ ChromaDB or Embedder not loaded. Cannot perform vector search.")
387
+ return [] # Return empty if Chroma or Embedder is not loaded
388
+
389
+ where_filter = {}
390
+ if model_filter:
391
+ where_filter["model"] = model_filter.lower()
392
+ if doc_type_filter:
393
+ where_filter["doc_type"] = doc_type_filter.lower()
394
+
395
+ # ChromaDB query expects a dictionary for 'where'
396
+ results = collection.query(
397
+ query_texts=[question],
398
+ n_results=top_k * 5, # fetch more for reranking
399
+ where={} if not where_filter else where_filter # Pass empty dict if no filter
400
+ )
401
+
402
+
403
+ if not results or not results.get("documents") or not results["documents"][0]:
404
+ return [] # No matches
405
+
406
+ try:
407
+ question_embedding = embedder.encode(question, convert_to_tensor=True)
408
+ except Exception as e:
409
+ print(f"Error encoding question: {e}")
410
+ return [] # Return empty if embedding fails
411
+
412
+ # Step 3: Compute semantic + keyword score
413
+ reranked = []
414
+ # Ensure results["documents"] and results["metadatas"] are not empty before iterating
415
+ if results.get("documents") and results["documents"][0]:
416
+ for i, text in enumerate(results["documents"][0]):
417
+ meta = results["metadatas"][0][i]
418
+
419
+ # Handle potential encoding errors during text embedding
420
+ try:
421
+ embedding = embedder.encode(text, convert_to_tensor=True)
422
+ # Semantic similarity
423
+ similarity_score = float(util.cos_sim(question_embedding, embedding))
424
+ except Exception as e:
425
+ print(f"Error encoding chunk text for reranking: {e}. Skipping chunk.")
426
+ continue # Skip this chunk if encoding fails
427
+
428
+
429
+ # Keyword score
430
+ keyword_score = 0
431
+ if rerank_keywords and text: # Ensure text is not None or empty
432
+ for kw in rerank_keywords:
433
+ if kw.lower() in text.lower():
434
+ keyword_score += 1
435
+
436
+ # Combine with tunable weights
437
+ # Weights should sum to 1 for a simple weighted average
438
+ final_score = (0.8 * similarity_score) + (0.2 * keyword_score)
439
+
440
+ reranked.append({
441
+ "score": final_score,
442
+ "text": text,
443
+ "metadata": meta
444
+ })
445
+
446
+ # Sort by combined score
447
+ reranked.sort(key=lambda x: x["score"], reverse=True)
448
+ return reranked[:top_k]
449
+
450
+
451
+ # ---------------------------
452
+ # 💬 Ask Hugging Face Model
453
+ # ---------------------------
454
+ def ask_hf_model(prompt):
455
+ if pipe is None:
456
+ return "Hugging Face model not loaded. Cannot generate response."
457
+ try:
458
+ # Use the Llama 3.1 chat template
459
+ messages = [
460
+ {"role": "system", "content": "You are a technical assistant trained to answer questions using equipment manuals. Use only the provided context to answer the question. If the answer is not clearly in the context, reply: 'I don't know.'"},
461
+ {"role": "user", "content": prompt}
462
+ ]
463
+
464
+ # Apply chat template and generate text
465
+ prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
466
+
467
+ outputs = pipe(
468
+ prompt_text,
469
+ do_sample=True,
470
+ temperature=0.1, # Keep temperature low for more factual answers
471
+ top_p=0.9,
472
+ max_new_tokens=512,
473
+ pad_token_id=tokenizer.eos_token_id # Set pad_token_id for generation
474
+ )
475
+ # The output includes the prompt, we need to extract just the generated part
476
+ # Find the end of the prompt text in the generated output
477
+ generated_text = outputs[0]["generated_text"]
478
+ # The chat template adds the assistant's turn token, look for that to find the response start
479
+ response_start_token = tokenizer.apply_chat_template([{"role": "assistant", "content": ""}], tokenize=False, add_generation_prompt=False)
480
+ response_start_index = generated_text.find(response_start_token)
481
+
482
+ if response_start_index != -1:
483
+ response = generated_text[response_start_index + len(response_start_token):].strip()
484
+ else:
485
+ # Fallback if the assistant token isn't found
486
+ response = generated_text.strip()
487
+
488
+ # Remove any trailing EOS tokens or similar artifacts
489
+ if response.endswith(tokenizer.eos_token):
490
+ response = response[:-len(tokenizer.eos_token)].strip()
491
+
492
+
493
+ return response
494
+
495
+ except Exception as e:
496
+ return f"❌ Error generating response from Hugging Face model: {str(e)}"
497
+
498
+ # ---------------------------
499
+ # 🎯 Full RAG Pipeline
500
+ # ---------------------------
501
+ def run_rag_qa(user_question, model_filter=None, doc_type_filter=None): # Added filters as optional inputs
502
+ # Ensure ChromaDB and the HF model pipeline are loaded before proceeding
503
+ if collection is None:
504
+ return "ChromaDB is not loaded. Ensure PDFs are in ./Manuals and the app started correctly."
505
+ if pipe is None:
506
+ return "Hugging Face model pipeline is not loaded. Ensure HF_TOKEN is set and the model loaded successfully."
507
+
508
+ results = query_manuals(
509
+ question=user_question,
510
+ model_filter=model_filter, # Use the optional filter inputs
511
+ doc_type_filter=doc_type_filter,
512
+ top_k=MAX_CONTEXT_CHUNKS,
513
+ rerank_keywords=["diagnostic", "immobilize", "system", "screen", "service", "error"] # Example keywords
514
+ )
515
+
516
+ if not results:
517
+ # Attempt a broader search if initial filter yields no results
518
+ if model_filter or doc_type_filter:
519
+ print("No results with specified filters, trying broader search...")
520
+ results = query_manuals(
521
+ question=user_question,
522
+ model_filter=None, # Remove filters for broader search
523
+ doc_type_filter=None,
524
+ top_k=MAX_CONTEXT_CHUNKS,
525
+ rerank_keywords=["diagnostic", "immobilize", "system", "screen", "service", "error"]
526
+ )
527
+ if not results:
528
+ return "No relevant documents found for the query, even with broader search."
529
+ else:
530
+ return "No relevant documents found for the query."
531
+
532
+
533
+ context = "\n\n".join([f"Source File: {r['metadata'].get('source_file', 'N/A')}, Page: {r['metadata'].get('page', 'N/A')}\nText: {r['text'].strip()}" for r in results])
534
+
535
+ prompt = f"""
536
+ Context:
537
+ {context}
538
+
539
+ Question: {user_question}
540
  """
541
+
542
+ return ask_hf_model(prompt)
543
+
544
+ # ---------------------------
545
+ # --- Initial Setup ---
546
+ # This code runs when the app starts on Hugging Face Spaces
547
+ # It processes PDFs, chunks, and builds the ChromaDB
548
+ # ---------------------------
549
+
550
+ print("Starting initial setup...")
551
+
552
+ # Ensure Tesseract is available on the system (Hugging Face Spaces usually has it, but this command is good practice)
553
+ # Using ! in app.py is generally discouraged, better to ensure the environment has it
554
+ # For HF Spaces, you might need to use a Dockerfile or rely on the default environment.
555
+ # If Tesseract isn't found, the OCR part might fail.
556
+
557
+ # Process PDFs and extract pages
558
+ all_pages = process_pdfs_for_pages(pdf_folder, output_jsonl_pages)
559
+
560
+ # Chunk the pages
561
+ all_chunks = []
562
+ if all_pages: # Only chunk if pages were processed
563
+ all_chunks = chunk_pages(output_jsonl_pages, output_jsonl_chunks, chunk_size, chunk_overlap)
564
+
565
+ # Embed chunks into ChromaDB
566
+ collection = None # Initialize collection
567
+ if all_chunks: # Only embed if chunks were created
568
+ collection, embed_error = embed_chunks_into_chroma(output_jsonl_chunks, chroma_path, collection_name)
569
+ if embed_error:
570
+ print(f"Error during embedding: {embed_error}")
571
+
572
+
573
+ print("Initial setup complete.")
574
+
575
+ # ---------------------------
576
+ # 🖥️ Gradio Interface
577
+ # ---------------------------
578
+ # Only define and launch the interface if the necessary components loaded
579
+ if collection is not None and pipe is not None:
580
+ with gr.Blocks() as demo:
581
+ gr.Markdown("""# 🧠 Manual QA via Hugging Face Llama 3.1
582
+ Ask a technical question and get answers using your own PDF manual database and a Hugging Face model.
583
+ **Note:** Initial startup might take time to process manuals and build the search index. Ensure your `Manuals` folder is uploaded and the `HF_TOKEN` secret is set in Space settings.
584
+ """)
585
+ with gr.Row():
586
+ question = gr.Textbox(label="Your Question", placeholder="e.g. How do I access diagnostics on the SE3 console?")
587
+ with gr.Row():
588
+ model_filter_input = gr.Textbox(label="Filter by Model (Optional)", placeholder="e.g. se3hd")
589
+ doc_type_filter_input = gr.Dropdown(label="Filter by Document Type (Optional)", choices=["owner manual", "service manual", "assembly instructions", "installer alert", "parts manual", "service bulletin", "unknown", None], value=None, allow_custom_value=True)
590
+
591
+ submit = gr.Button("🔍 Ask")
592
+ answer = gr.Textbox(label="Answer", lines=10) # Increased lines for better readability
593
+
594
+ # Call the run_rag_qa function when the button is clicked
595
+ submit.click(
596
+ fn=run_rag_qa,
597
+ inputs=[question, model_filter_input, doc_type_filter_input],
598
+ outputs=[answer]
599
+ )
600
+
601
+ # In Hugging Face Spaces, the app is launched automatically.
602
+ # The demo.launch() call is removed.
603
+ # demo.launch()
604
+ else:
605
+ print("Gradio demo will not launch because RAG components (ChromaDB or HF Model) failed to load during setup.")
606
+ # You could add a simple Gradio interface here to show an error message
607
+ # if you wanted to provide user feedback in the Space UI even on failure.
608
+ # Example:
609
+ # with gr.Blocks() as error_demo:
610
+ # gr.Markdown("## Application Failed to Load")
611
+ # gr.Textbox(label="Error Details", value="RAG components (ChromaDB or HF Model) failed to initialize. Check logs and Space settings (HF_TOKEN, resources).", interactive=False)
612
+ # error_demo.launch()