Spaces:
Sleeping
Sleeping
File size: 3,797 Bytes
30e8063 a3cddef 30e8063 a3cddef 30e8063 a3cddef 30e8063 a3cddef 30e8063 a3cddef 30e8063 a3cddef 30e8063 a3cddef 30e8063 a3cddef 30e8063 | 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 108 109 110 111 112 113 114 | import streamlit as st
from sentence_transformers import SentenceTransformer, util
import PyPDF2
from PIL import Image
import pytesseract
import speech_recognition as sr
import docx
import tempfile
import os
# Load embedding model
model = SentenceTransformer('all-MiniLM-L6-v2')
# Database to store file content and embeddings
file_db = []
# Function to extract text from a PDF
def extract_text_from_pdf(pdf_file):
pdf_reader = PyPDF2.PdfReader(pdf_file)
text = ""
for page in pdf_reader.pages:
text += page.extract_text()
return text
# Function to extract text from an image
def extract_text_from_image(image_file):
image = Image.open(image_file)
text = pytesseract.image_to_string(image)
return text
# Function to extract text from a video (audio)
def extract_text_from_video(video_file):
r = sr.Recognizer()
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
temp_audio_path = temp_audio.name
os.system(f"ffmpeg -i {video_file.name} -vn -ar 16000 -ac 1 -f wav {temp_audio_path}")
with sr.AudioFile(temp_audio_path) as source:
audio_data = r.record(source)
os.remove(temp_audio_path)
return r.recognize_google(audio_data)
# Function to extract text from a TXT file
def extract_text_from_txt(txt_file):
return txt_file.read().decode("utf-8")
# Function to extract text from a Word document
def extract_text_from_word(docx_file):
doc = docx.Document(docx_file)
text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
return text
# Function to process files and generate embeddings
def process_file(file):
file_type = file.type.split('/')[1]
content = ""
try:
if file.type == 'application/pdf': # PDF files
content = extract_text_from_pdf(file)
elif file.type in ['image/png', 'image/jpeg']: # Images
content = extract_text_from_image(file)
elif file.type == 'video/mp4': # Videos
content = extract_text_from_video(file)
elif file.type == 'text/plain': # TXT files
content = extract_text_from_txt(file)
elif file.name.endswith('.docx'): # Word documents
content = extract_text_from_word(file)
except Exception as e:
st.error(f"Error processing file {file.name}: {e}")
if content:
embedding = model.encode(content, convert_to_tensor=True)
file_db.append({"name": file.name, "content": content, "embedding": embedding})
st.success(f"File '{file.name}' processed successfully!")
else:
st.warning(f"Could not extract text from '{file.name}'")
# Streamlit UI
st.title("AI-based Text Search for Files")
st.write("Upload documents, images, or videos, and perform AI-powered text search.")
# File upload section
uploaded_files = st.file_uploader(
"Upload your files",
type=['pdf', 'png', 'jpg', 'jpeg', 'mp4', 'txt', 'docx'],
accept_multiple_files=True
)
if uploaded_files:
for file in uploaded_files:
process_file(file)
# Search section
query = st.text_input("Enter your search query:")
if query:
query_embedding = model.encode(query, convert_to_tensor=True)
results = []
for file in file_db:
similarity = util.pytorch_cos_sim(query_embedding, file['embedding'])[0][0].item()
results.append({"name": file['name'], "similarity": similarity, "content": file['content']})
# Sort results by similarity
results = sorted(results, key=lambda x: x['similarity'], reverse=True)
st.write("Search Results:")
for result in results[:5]: # Display top 5 results
st.write(f"**{result['name']}** (Similarity: {result['similarity']:.2f})")
st.write(result['content'][:500] + "...") # Show a preview |