Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import os | |
| from langdetect import detect | |
| from PyPDF2 import PdfReader | |
| import requests | |
| from sentence_transformers import SentenceTransformer | |
| import faiss | |
| import numpy as np | |
| # Load the API key from Streamlit secrets | |
| API_KEY = st.secrets["Key2"] | |
| API_URL = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-alpha" | |
| # Load the embedding model for semantic search | |
| embedding_model = SentenceTransformer('all-MiniLM-L6-v2') | |
| # Function to query the LLM via Hugging Face Inference API | |
| def query_llm_api(prompt, max_new_tokens=1000, temperature=0.7, top_k=50): | |
| headers = { | |
| "Authorization": f"Bearer {API_KEY}", | |
| "Content-Type": "application/json", | |
| } | |
| payload = { | |
| "inputs": prompt, | |
| "parameters": { | |
| "max_new_tokens": max_new_tokens, | |
| "temperature": temperature, | |
| "top_k": top_k, | |
| }, | |
| } | |
| response = requests.post(API_URL, headers=headers, json=payload) | |
| if response.status_code == 200: | |
| return response.json()["generated_text"] | |
| else: | |
| st.error(f"Error querying the API: {response.status_code}, {response.text}") | |
| return None | |
| # Function to detect language | |
| def detect_language(text): | |
| try: | |
| return detect(text) | |
| except Exception: | |
| return "en" # Default to English if detection fails | |
| # Function to extract text from PDF with line and page numbers | |
| def extract_text_from_pdf(pdf_file): | |
| pdf_reader = PdfReader(pdf_file) | |
| text_data = [] | |
| for page_num, page in enumerate(pdf_reader.pages): | |
| if page.extract_text(): | |
| lines = page.extract_text().split('\n') | |
| for line_num, line in enumerate(lines): | |
| text_data.append({ | |
| "page": page_num + 1, | |
| "line": line_num + 1, | |
| "content": line | |
| }) | |
| return text_data | |
| # Function to create embeddings for the PDF text | |
| def get_embeddings(text_data): | |
| texts = [entry['content'] for entry in text_data] | |
| return embedding_model.encode(texts, convert_to_tensor=False) | |
| # Function to perform KNN or cosine similarity search | |
| def search_pdf_content(pdf_text_data, query, search_type="knn", k=5): | |
| query_embedding = embedding_model.encode([query])[0] | |
| pdf_embeddings = get_embeddings(pdf_text_data) | |
| if search_type == "knn": | |
| index = faiss.IndexFlatL2(pdf_embeddings.shape[1]) | |
| index.add(pdf_embeddings.astype('float32')) | |
| distances, indices = index.search(np.array([query_embedding], dtype='float32'), k) | |
| return [pdf_text_data[i] for i in indices[0]] | |
| elif search_type == "cosine": | |
| pdf_embeddings_norm = pdf_embeddings / np.linalg.norm(pdf_embeddings, axis=1, keepdims=True) | |
| query_embedding_norm = query_embedding / np.linalg.norm(query_embedding) | |
| similarities = np.dot(pdf_embeddings_norm, query_embedding_norm) | |
| top_indices = np.argsort(similarities)[-k:][::-1] | |
| return [pdf_text_data[i] for i in top_indices] | |
| # Streamlit UI | |
| st.title("PDF Search with LLM and Semantic Search") | |
| pdf_file = st.file_uploader("Upload a PDF file", type="pdf") | |
| search_query = st.text_input("Enter your search query") | |
| search_method = st.radio("Select Search Method", ("knn", "cosine")) | |
| k_value = st.slider("Number of Results (K)", min_value=1, max_value=20, value=5) | |
| if pdf_file and search_query: | |
| pdf_text_data = extract_text_from_pdf(pdf_file) | |
| results = search_pdf_content(pdf_text_data, search_query, search_type=search_method, k=k_value) | |
| st.write("### Search Results") | |
| for res in results: | |
| st.write(f"**Page {res['page']}, Line {res['line']}:** {res['content']}") | |