| print("start1") |
| import os |
| import sys |
| import subprocess |
| import gradio as gr |
| from PyPDF2 import PdfReader |
| from langchain.text_splitter import RecursiveCharacterTextSplitter |
| from langchain_community.vectorstores import FAISS |
| from langchain.prompts import PromptTemplate |
| from langchain.chains import LLMChain |
| from langchain_community.embeddings import HuggingFaceEmbeddings |
| from langchain.schema import Document |
| print("start2") |
|
|
|
|
| # Check if already installed to avoid reinstalling |
| try: |
| import llama_cpp |
| print("llama_cpp already installed.") |
| except ImportError: |
| print("Installing llama_cpp from wheel...") |
| subprocess.check_call([ |
| sys.executable, "-m", "pip", "install", |
| "llama-cpp-python", "--no-binary", ":all:", "--force-reinstall" |
| ]) |
|
|
|
|
| from llama_cpp import Llama |
| print("start3") |
| import warnings |
| warnings.filterwarnings("ignore") |
|
|
| print("Start") |
| import subprocess |
|
|
| subprocess.run([ |
| "huggingface-cli", "download", |
| "TheBloke/Mistral-7B-Instruct-v0.1-GGUF", |
| "mistral-7b-instruct-v0.1.Q2_K.gguf", |
| "--local-dir", "./models", |
| "--local-dir-use-symlinks", "False" |
| ], check=True) |
|
|
| # ------------------------------ |
| # Device and Embedding Setup (CPU optimized) |
| # ------------------------------ |
| modelPath = "sentence-transformers/all-mpnet-base-v2" |
| model_kwargs = {"device": "cpu"} # Force CPU usage |
| encode_kwargs = {"normalize_embedding": False} |
|
|
| embeddings = HuggingFaceEmbeddings( |
| model_name=modelPath, |
| model_kwargs=model_kwargs, |
| encode_kwargs=encode_kwargs |
| ) |
|
|
| # ------------------------------ |
| # Load Mistral GGUF via llama.cpp (CPU optimized) |
| # ------------------------------ |
| llm_cpp = Llama( |
| model_path="./models/mistral-7b-instruct-v0.1.Q2_K.gguf", |
| n_ctx=2048, |
| n_threads=4, # Adjust based on your CPU cores |
| n_gpu_layers=0, # Force CPU-only |
| temperature=0.7, |
| top_p=0.9, |
| repeat_penalty=1.1 |
| ) |
|
|
| # ------------------------------ |
| # LangChain-compatible wrapper |
| # ------------------------------ |
| def mistral_llm(prompt): |
| output = llm_cpp( |
| prompt, |
| max_tokens=512, # Reduced for CPU performance |
| stop=["</s>", "[INST]", "[/INST]"] |
| ) |
| return output["choices"][0]["text"].strip() |
|
|
| # ------------------------------ |
| # Prompt Template (unchanged) |
| # ------------------------------ |
| def get_qa_prompt(): |
| template = """<s>[INST] \ |
| You are a helpful, knowledgeable AI assistant. Answer the user's question based on the provided context. |
|
|
| Guidelines: |
| - Respond in a natural, conversational tone |
| - Be detailed but concise |
| - Use paragraphs and bullet points when appropriate |
| - If you don't know, say so |
| - Maintain a friendly and professional demeanor |
|
|
| Conversation History: |
| {chat_history} |
|
|
| Relevant Context: |
| {context} |
|
|
| Current Question: {question} |
|
|
| Provide a helpful response: [/INST]""" |
| return PromptTemplate( |
| template=template, |
| input_variables=["context", "question", "chat_history"] |
| ) |
|
|
| # ------------------------------ |
| # PDF and Chat Logic (optimized for CPU) |
| # ------------------------------ |
| def pdf_text(pdf_docs): |
| text = "" |
| for doc in pdf_docs: |
| reader = PdfReader(doc) |
| for page in reader.pages: |
| page_text = page.extract_text() |
| if page_text: |
| text += page_text + "\n" |
| return text |
|
|
| def get_chunks(text): |
| splitter = RecursiveCharacterTextSplitter( |
| chunk_size=800, # Smaller chunks for CPU |
| chunk_overlap=100, |
| length_function=len |
| ) |
| chunks = splitter.split_text(text) |
| return [Document(page_content=chunk) for chunk in chunks] |
|
|
| def get_vectorstore(documents): |
| db = FAISS.from_documents(documents, embedding=embeddings) |
| db.save_local("faiss_index") |
|
|
| def format_chat_history(history): |
| return "\n".join([f"User: {q}\nAssistant: {a}" for q, a in history[-2:]]) # Shorter history |
|
|
| def handle_pdf_upload(pdf_files): |
| if not pdf_files: |
| return "⚠️ Upload at least one PDF" |
| try: |
| text = pdf_text(pdf_files) |
| if not text.strip(): |
| return "⚠️ Could not extract text" |
| chunks = get_chunks(text) |
| get_vectorstore(chunks) |
| return f"✅ Processed {len(pdf_files)} PDF(s) with {len(chunks)} chunks" |
| except Exception as e: |
| return f"❌ Error: {str(e)}" |
|
|
| def user_query(msg, chat_history): |
| if not os.path.exists("faiss_index"): |
| chat_history.append((msg, "Please upload PDF documents first.")) |
| return "", chat_history |
|
|
| try: |
| db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True) |
| retriever = db.as_retriever(search_kwargs={"k": 2}) # Fewer documents for CPU |
| docs = retriever.get_relevant_documents(msg) |
| context = "\n\n".join([d.page_content for d in docs][:2]) # Limit context |
|
|
| prompt = get_qa_prompt() |
| final_prompt = prompt.format( |
| context=context[:1500], # Further limit context size |
| question=msg, |
| chat_history=format_chat_history(chat_history) |
| ) |
|
|
| response = mistral_llm(final_prompt) |
| chat_history.append((msg, response)) |
| return "", chat_history |
| except Exception as e: |
| error_msg = f"Sorry, I encountered an error: {str(e)}" |
| chat_history.append((msg, error_msg)) |
| return "", chat_history |
|
|
| # ------------------------------ |
| # Gradio Interface (your exact requested format) |
| # ------------------------------ |
| with gr.Blocks(theme=gr.themes.Soft(), title="PDF Chat Assistant") as demo: |
| with gr.Row(): |
| gr.Markdown(""" |
| # 📚 PDF Chat Assistant |
| ### Have natural conversations with your documents ((Note: This Space runs on CPU, so responses may take a few mins.)) |
| """) |
| with gr.Row(): |
| with gr.Column(scale=1, min_width=300): |
| gr.Markdown("### Document Upload") |
| pdf_input = gr.File( |
| file_types=[".pdf"], |
| file_count="multiple", |
| label="Upload PDFs", |
| height=100 |
| ) |
| upload_btn = gr.Button("Process Documents", variant="primary") |
| status_box = gr.Textbox(label="Status", interactive=False) |
| gr.Markdown(""" |
| **Instructions:** |
| 1. Upload PDF documents |
| 2. Click Process Documents |
| 3. Start chatting in the right panel |
| """) |
|
|
| with gr.Column(scale=2): |
| chatbot = gr.Chatbot( |
| height=600, |
| bubble_full_width=False, |
| avatar_images=( |
| "user.png", |
| "bot.png" |
| ) |
| ) |
| |
| with gr.Row(): |
| message = gr.Textbox( |
| placeholder="Type your question about the documents...", |
| show_label=False, |
| container=False, |
| scale=7, |
| autofocus=True |
| ) |
| submit_btn = gr.Button("Send", variant="primary", scale=1) |
| |
| with gr.Row(): |
| clear_chat = gr.Button("🧹 Clear Conversation") |
| examples = gr.Examples( |
| examples=[ |
| "Summarize the key points from the documents", |
| "What are the main findings?", |
| "Explain this in simpler terms" |
| ], |
| inputs=message, |
| label="Example Questions" |
| ) |
|
|
| upload_btn.click(handle_pdf_upload, inputs=pdf_input, outputs=status_box) |
| submit_btn.click(user_query, inputs=[message, chatbot], outputs=[message, chatbot]) |
| message.submit(user_query, inputs=[message, chatbot], outputs=[message, chatbot]) |
| clear_chat.click(lambda: [], None, chatbot, queue=False) |
|
|
| if __name__ == "__main__": |
| demo.launch() # Disable sharing for local CPU use |
