Spaces:
Running
Running
| import gradio as gr | |
| from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline | |
| from pdfminer.high_level import extract_text | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| import numpy as np | |
| import nltk | |
| from nltk.tokenize import sent_tokenize | |
| from rank_bm25 import BM25Okapi | |
| nltk.download('punkt') | |
| # QA model | |
| qa_model_name = "deepset/roberta-large-squad2" | |
| qa_model = AutoModelForQuestionAnswering.from_pretrained(qa_model_name) | |
| qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name) | |
| qa_pipeline = pipeline("question-answering", model=qa_model, tokenizer=qa_tokenizer) | |
| # Summarization model | |
| summarization_model_name = "facebook/bart-large-cnn" | |
| summarizer = pipeline("summarization", model=summarization_model_name) | |
| def read_pdf(file): | |
| try: | |
| text = extract_text(file) | |
| if not text: | |
| raise ValueError("PDF extraction failed.") | |
| return text | |
| except Exception as e: | |
| return str(e) | |
| def retrieve_relevant_text_bm25(question, sentences, top_n=3): | |
| try: | |
| tokenized_corpus = [sent.split() for sent in sentences] | |
| bm25 = BM25Okapi(tokenized_corpus) | |
| tokenized_query = question.split() | |
| doc_scores = bm25.get_scores(tokenized_query) | |
| top_n_indices = np.argsort(doc_scores)[::-1][:top_n] | |
| relevant_texts = [sentences[i] for i in top_n_indices] | |
| return " ".join(relevant_texts) | |
| except Exception as e: | |
| return str(e) | |
| def answer_question(pdf, question, num_words): | |
| try: | |
| text = read_pdf(pdf) | |
| if isinstance(text, str): | |
| return text | |
| if "summarize" in question.lower(): | |
| try: | |
| summarized_text = summarizer(text, max_length=num_words, min_length=1) | |
| return summarized_text[0]['summary_text'].strip() | |
| except RuntimeError as e: | |
| if "Input length of input_ids is" in str(e) and "but `max_length` is set to" in str(e): | |
| return "PDF is too long for summarization. Please provide a shorter PDF or ask a more specific question." | |
| else: | |
| return f"Summarization Error: {e}" | |
| except Exception as e: | |
| return f"Summarization Error: {e}" | |
| sentences = sent_tokenize(text) | |
| relevant_text = retrieve_relevant_text_bm25(question, sentences) | |
| if not relevant_text: | |
| return "Could not find relevant information in the PDF." | |
| response = qa_pipeline(question=question, context=relevant_text) | |
| answer = response.get('answer') | |
| if not answer: | |
| return "Could not find an answer in the relevant text." | |
| answer = answer.strip() | |
| answer = " ".join(answer.split()) | |
| if len(answer.split()) > num_words: | |
| try: | |
| summarized_answer = summarizer(answer, max_length=num_words + 10, min_length=1) | |
| answer = summarized_answer[0]['summary_text'] | |
| answer = answer.strip() | |
| answer = " ".join(answer.split()) | |
| if len(answer.split()) > num_words: | |
| answer = " ".join(answer.split()[:num_words]) | |
| except RuntimeError as e: | |
| if "Input length of input_ids is" in str(e) and "but `max_length` is set to" in str(e): | |
| answer = " ".join(answer.split()[:num_words]) | |
| else: | |
| return f"Summarization Error: {e}" | |
| except Exception as e: | |
| return f"Summarization Error: {e}" | |
| elif len(answer.split()) < num_words and relevant_text: | |
| remaining_words = num_words - len(answer.split()) | |
| added_words = 0 | |
| added_sentences = [] | |
| for sentence in sent_tokenize(relevant_text): | |
| sentence_words = sentence.split() | |
| words_to_add = min(remaining_words - added_words, len(sentence_words)) | |
| if words_to_add > 0: | |
| added_sentences.append(" ".join(sentence_words[:words_to_add])) | |
| added_words += words_to_add | |
| if added_words >= remaining_words: | |
| break | |
| answer += " " + " ".join(added_sentences) | |
| answer = answer.strip() | |
| answer = " ".join(answer.split()) | |
| if len(answer.split()) > num_words: | |
| answer = " ".join(answer.split()[:num_words]) | |
| return answer.strip() | |
| except Exception as e: | |
| return str(e) | |
| with gr.Blocks() as iface: | |
| gr.Markdown("PDF Q&A with RoBERTa | made by NP") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| question_input = gr.Textbox(lines=2, placeholder="Ask a question", label="Question") | |
| btn = gr.Button("Submit") | |
| with gr.Column(scale=1): | |
| pdf_input = gr.File(type="filepath", label="Upload PDF") | |
| num_words_slider = gr.Slider(minimum=1, maximum=500, value=100, step=1, label="Number of Words") | |
| answer_output = gr.Textbox(label="Answer", lines=5) | |
| btn.click(fn=answer_question, inputs=[pdf_input, question_input, num_words_slider], outputs=answer_output) | |
| if __name__ == "__main__": | |
| iface.launch() |