Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| from PIL import Image | |
| import pytesseract | |
| from pdf2image import convert_from_path | |
| from langchain_community.embeddings import HuggingFaceEmbeddings | |
| from langchain.prompts import PromptTemplate | |
| from langchain.chains import RetrievalQA | |
| from langchain.memory import ConversationBufferMemory | |
| from langchain_groq import ChatGroq | |
| from langchain_community.vectorstores import FAISS | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| # Initialize the Groq API Key and the model | |
| os.environ["GROQ_API_KEY"] = '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' | |
| llm = ChatGroq( | |
| model='llama3-70b-8192', | |
| temperature=0.5, | |
| max_tokens=None, | |
| timeout=None, | |
| max_retries=2 | |
| ) | |
| # OCR functions | |
| def ocr_image(image_path, language='eng+guj'): | |
| img = Image.open(image_path) | |
| return pytesseract.image_to_string(img, lang=language) | |
| def ocr_pdf(pdf_path, language='eng+guj'): | |
| images = convert_from_path(pdf_path) | |
| all_text = "\n".join(pytesseract.image_to_string(img, lang=language) for img in images) | |
| return all_text | |
| def ocr_file(file_path): | |
| ext = os.path.splitext(file_path)[1].lower() | |
| if ext == ".pdf": | |
| return ocr_pdf(file_path) | |
| elif ext in [".jpg", ".jpeg", ".png", ".bmp"]: | |
| return ocr_image(file_path) | |
| else: | |
| return "Unsupported file format." | |
| def get_text_chunks(text): | |
| splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100) | |
| return splitter.split_text(text) | |
| def get_vector_store(chunks): | |
| embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") | |
| vector_store = FAISS.from_texts(chunks, embedding=embeddings) | |
| os.makedirs("faiss_index", exist_ok=True) | |
| vector_store.save_local("faiss_index") | |
| return vector_store | |
| # Conversational chain | |
| def get_conversational_chain(): | |
| template = """<Insert your prompt here>""" | |
| embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-MiniLM-L6-v2") | |
| vector_store = FAISS.load_local("faiss_index", embeddings) | |
| qa_chain = RetrievalQA.from_chain_type( | |
| llm, | |
| retriever=vector_store.as_retriever(), | |
| chain_type='stuff', | |
| verbose=True, | |
| chain_type_kwargs={ | |
| "prompt": PromptTemplate(input_variables=["history", "context", "question"], template=template), | |
| "memory": ConversationBufferMemory(memory_key="history", input_key="question"), | |
| } | |
| ) | |
| return qa_chain | |
| # File and question handling | |
| def process_files(files, question): | |
| text = "" | |
| for file in files: | |
| file_path = os.path.join("temp", file.name) | |
| with open(file_path, "wb") as f: | |
| f.write(file.read()) | |
| text += ocr_file(file_path) + "\n" | |
| chunks = get_text_chunks(text) | |
| vector_store = get_vector_store(chunks) | |
| qa_chain = get_conversational_chain() | |
| response = qa_chain({"query": question}) | |
| return response.get("result", "No result found.") | |
| # Gradio Interface | |
| def app(files, question): | |
| return process_files(files, question) | |
| iface = gr.Interface( | |
| fn=app, | |
| inputs=[gr.File(file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp"], label="Upload Files"), gr.Textbox(label="Ask a Question")], | |
| outputs="text", | |
| title="OCR and Document Query System", | |
| description="Upload PDF or image files and ask questions based on their content." | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() | |