# app.py import os import re import numpy as np import torch from sentence_transformers import SentenceTransformer import pdfplumber from chromadb import Client from chromadb.config import Settings from chromadb.utils import embedding_functions from transformers import AutoTokenizer from rerankers import Reranker from transformers import GPT2TokenizerFast from groq import Groq from chromadb import PersistentClient import gradio as gr # Retrieve the API key from environment variables (Hugging Face Secrets) groq_api_key = os.environ.get('GROQ_API_KEY') # Initialize the chat client with the API key chat_client = Groq(api_key=groq_api_key) model = "llama-3.2-90b-text-preview" def parse_pdf(pdf_path): texts = [] with pdfplumber.open(pdf_path) as pdf: for page_num, page in enumerate(pdf.pages, start=1): text = page.extract_text() if text: texts.append({ 'text': text, 'metadata': { 'page_number': page_num } }) return texts def preprocess_text(text): # ... (same as your original function) text = re.sub(r'\s+', ' ', text) text = text.strip() return text def call_Llama_api(query, context): # ... (same as your original function) chat_completion = chat_client.chat.completions.create( messages=[ { "role": "system", "content": "You are a car technician. Given the user's question and relevant excerpts from different car manuals, answer the question by including direct quotes from the correct car manual. Be concise and to the point in your response." }, { "role": "user", "content": "User Question: " + query + "\n\nRelevant Excerpt(s):\n\n" + context, } ], temperature=0.7, max_tokens=50, top_p=1, stream=False, stop=None, model=model ) response = chat_completion.choices[0].message.content return response def chunk_texts(texts, max_tokens=500, overlap_tokens=50): """ Splits texts into chunks based on paragraphs with overlap to preserve context. """ chunks = [] for item in texts: text = preprocess_text(item['text']) if not text: continue metadata = item['metadata'] # Split text into paragraphs paragraphs = text.split('\n\n') current_chunk = '' current_tokens = 0 for i, paragraph in enumerate(paragraphs): paragraph = paragraph.strip() if not paragraph: continue paragraph_tokens = len(tokenizer.encode(paragraph)) if current_tokens + paragraph_tokens <= max_tokens: current_chunk += paragraph + '\n\n' current_tokens += paragraph_tokens else: # Save the current chunk chunk = { 'text': current_chunk.strip(), 'metadata': metadata } chunks.append(chunk) # Start a new chunk with overlap overlap_text = ' '.join(current_chunk.split()[-overlap_tokens:]) current_chunk = overlap_text + ' ' + paragraph + '\n\n' current_tokens = len(tokenizer.encode(current_chunk)) if current_chunk: chunk = { 'text': current_chunk.strip(), 'metadata': metadata } chunks.append(chunk) return chunks def is_car_model_available(query, available_models): # ... (same as your original function) for model in available_models: if model.lower() in query.lower(): return model return None def extract_car_model(pdf_filename): base_name = os.path.basename(pdf_filename) match = re.search(r'manual_(.+)\.pdf', base_name) if match: model_name = match.group(1).replace('_', ' ').title() return model_name else: return 'Unknown Model' def colbert_rerank(query=None, chunks=None): # ... (same as your original function) d = ranker.rank(query=query, docs=chunks) reranked_chunks = [d[i].text for i in range(len(chunks))] return reranked_chunks[:10] def process_query(query): # Use global variables global available_car_models, collection car_model = is_car_model_available(query, available_car_models) if not car_model: return "The manual for the specified car model is not present." # Initial retrieval from ChromaDB results = collection.query( query_texts=[query], n_results=50, where={"car_model": car_model}, include=['documents', 'metadatas'] ) if not results['documents']: return "No relevant information found in the manual." # Extract chunks and metadata chunks = results['documents'][0] metadatas = results['metadatas'][0] reranked_chunks = colbert_rerank(query, chunks) final_context = " ".join(reranked_chunks[:10]) answer = call_Llama_api(query, final_context) # Prepare citations citations = [ f"Page {meta.get('page_number', 'N/A')}" for meta in metadatas[:5] ] citations_text = "Citations:\n" + "\n".join(citations) return f"{answer}\n\n{citations_text}" # Initialize global variables def initialize(): global collection, available_car_models, ranker # Check for CUDA availability device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Using device: {device}") tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") # For token counting # Initialize embedding model embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction( model_name="all-MiniLM-L12-v2", device=device ) client = PersistentClient(path="./chromadb") # Get the collection collection_name = "car_manuals5" if collection_name in [col.name for col in client.list_collections()]: collection = client.get_collection( name=collection_name, embedding_function=embedding_function ) available_car_models = ['Tiago', 'Astor'] else: collection = client.create_collection( name=collection_name, embedding_function=embedding_function ) # collection = client.get_or_create_collection( # name=collection_name, # embedding_function=embedding_function # ) # Set available car models # available_car_models = ['TIAGO', 'Astor'] pdf_files = ['./car_manuals/manual_Tiago.pdf', './car_manuals/manual_Astor.pdf'] available_car_models = [] for pdf_file in pdf_files: print(f"Parsing {pdf_file}...") pdf_texts = parse_pdf(pdf_file) car_model = extract_car_model(pdf_file) available_car_models.append(car_model) # Add car model to metadata for item in pdf_texts: item['metadata']['car_model'] = car_model # Chunk texts using the refined strategy chunks = chunk_texts(pdf_texts, max_tokens=500, overlap_tokens=50) # Prepare data for ChromaDB documents = [chunk['text'] for chunk in chunks] metadatas = [chunk['metadata'] for chunk in chunks] ids = [f"{car_model}_{i}" for i in range(len(documents))] # Add to ChromaDB collection collection.add( documents=documents, metadatas=metadatas, ids=ids ) # Initialize the ranker ranker = Reranker("answerdotai/answerai-colbert-small-v1", model_type='colbert') # Call initialize function initialize() # Set up the Gradio interface iface = gr.Interface( fn=process_query, inputs=gr.inputs.Textbox(lines=2, placeholder='Enter your question here...'), outputs='text', title='Car Manual Assistant', description='Ask a question about your car manual.', ) if __name__ == "__main__": # iface.launch(server_name="0.0.0.0", server_port=7860) iface.launch()