Spaces:
Running
Running
File size: 3,506 Bytes
21753a3 d2c6ac6 b14a2f9 21753a3 b14a2f9 1fda785 b14a2f9 d2c6ac6 b14a2f9 21753a3 d2c6ac6 21753a3 94f2884 21753a3 d2c6ac6 21753a3 d2c6ac6 21753a3 94f2884 21753a3 d2c6ac6 21753a3 d2c6ac6 21753a3 1fda785 21753a3 1fda785 21753a3 1fda785 21753a3 d2c6ac6 21753a3 fb8d4f3 21753a3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
from flask import Flask, request, jsonify, render_template
import fitz # PyMuPDF for PDF text extraction
import faiss # FAISS for vector search
import numpy as np
from sentence_transformers import SentenceTransformer
from huggingface_hub import InferenceClient
from typing import List, Tuple
app = Flask(__name__)
# Default settings
class ChatConfig:
MODEL = "google/gemma-3-27b-it"
DEFAULT_SYSTEM_MSG = "You are an AI assistant answering only based on the uploaded PDF."
DEFAULT_MAX_TOKENS = 512
DEFAULT_TEMP = 0.3
DEFAULT_TOP_P = 0.95
client = InferenceClient(ChatConfig.MODEL)
embed_model = SentenceTransformer("all-MiniLM-L6-v2") # Lightweight embedding model
vector_dim = 384 # Embedding size
index = faiss.IndexFlatL2(vector_dim) # FAISS index
documents = [] # Store extracted text
def extract_text_from_pdf(pdf_path):
"""Extracts text from PDF"""
doc = fitz.open(pdf_path)
text_chunks = [page.get_text("text") for page in doc]
return text_chunks
def create_vector_db(text_chunks):
"""Embeds text chunks and adds them to FAISS index"""
global documents, index
documents = text_chunks
embeddings = embed_model.encode(text_chunks)
index.add(np.array(embeddings, dtype=np.float32))
def search_relevant_text(query):
"""Finds the most relevant text chunk for the given query"""
query_embedding = embed_model.encode([query])
_, closest_idx = index.search(np.array(query_embedding, dtype=np.float32), k=3)
return "\n".join([documents[i] for i in closest_idx[0]])
def generate_response(
message: str,
history: List[Tuple[str, str]],
system_message: str = ChatConfig.DEFAULT_SYSTEM_MSG,
max_tokens: int = ChatConfig.DEFAULT_MAX_TOKENS,
temperature: float = ChatConfig.DEFAULT_TEMP,
top_p: float = ChatConfig.DEFAULT_TOP_P
) -> str:
if not documents:
return "Please upload a PDF first."
context = search_relevant_text(message) # Get relevant content from PDF
messages = [{"role": "system", "content": system_message}]
for user_msg, bot_msg in history:
if user_msg:
messages.append({"role": "user", "content": user_msg})
if bot_msg:
messages.append({"role": "assistant", "content": bot_msg})
messages.append({"role": "user", "content": f"Context: {context}\nQuestion: {message}"})
response = ""
for chunk in client.chat_completion(
messages,
max_tokens=max_tokens,
stream=True,
temperature=temperature,
top_p=top_p,
):
token = chunk.choices[0].delta.content or ""
response += token
return response
@app.route('/')
def index():
"""Serve the HTML page for the user interface"""
return render_template('index.html')
@app.route('/upload_pdf', methods=['POST'])
def upload_pdf():
"""Handle PDF upload"""
file = request.files['pdf']
pdf_path = f"uploaded_files/{file.filename}"
file.save(pdf_path)
# Extract text and create vector database
text_chunks = extract_text_from_pdf(pdf_path)
create_vector_db(text_chunks)
return jsonify({"message": "PDF uploaded and indexed successfully!"})
@app.route('/ask_question', methods=['POST'])
def ask_question():
"""Handle user question"""
message = request.json.get('message')
history = request.json.get('history', [])
response = generate_response(message, history)
return jsonify({"response": response})
if __name__ == '__main__':
app.run(debug=True)
|