Update app.py

app.py

@@ -1,42 +1,63 @@
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from transformers.pipelines import pipeline
+from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
+from transformers import BigBirdTokenizer, BigBirdForSequenceClassification
 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
 
-# Load T5 model and tokenizer
-model_name = "t5-large"
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-qa_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+# Load BigBird for text extraction
+bb_model_name = "google/bigbird-roberta-base"
+bb_model = BigBirdForSequenceClassification.from_pretrained(bb_model_name)
+bb_tokenizer = BigBirdTokenizer.from_pretrained(bb_model_name)
+
+# Load BERT for QA
+bert_model_name = "bert-large-uncased-whole-word-masking-finetuned-squad"
+bert_model = AutoModelForQuestionAnswering.from_pretrained(bert_model_name)
+bert_tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
+qa_pipeline = pipeline("question-answering", model=bert_model, tokenizer=bert_tokenizer)
 
 def read_pdf(file):
-    return extract_text(file)
+    try:
+        text = extract_text(file)
+        if not text:
+            raise ValueError("PDF extraction failed. The PDF might be scanned or have an unsupported format.")
+        return text
+    except Exception as e:
+        return str(e)
 
 def retrieve_relevant_text(question, context, top_n=5):
-    vectorizer = TfidfVectorizer().fit_transform([question] + context)
-    vectors = vectorizer.toarray()
-    cosine_matrix = cosine_similarity(vectors)
-    similar_ix = np.argsort(cosine_matrix[0])[::-1][1:top_n+1]
-    relevant_texts = [context[ix] for ix in similar_ix]
-    return " ".join(relevant_texts)
+    try:
+        vectorizer = TfidfVectorizer().fit_transform([question] + context)
+        vectors = vectorizer.toarray()
+        cosine_matrix = cosine_similarity(vectors)
+        similar_ix = np.argsort(cosine_matrix[0])[::-1][1:top_n+1]
+        relevant_texts = [context[ix] for ix in similar_ix]
+        return " ".join(relevant_texts)
+    except Exception as e:
+        return str(e)
 
 def answer_question(pdf, question):
-    context = read_pdf(pdf).split("\n")
-    relevant_text = retrieve_relevant_text(question, context)
-    input_text = f"question: {question} context: {relevant_text}"
-    response = qa_pipeline(input_text, max_length=512, do_sample=False)
-    return response[0]['generated_text']
+    try:
+        context = read_pdf(pdf).split("\n")
+        if isinstance(context, str):
+            return context  # Return error message if read_pdf failed
+        relevant_text = retrieve_relevant_text(question, context)
+        if isinstance(relevant_text, str):
+            return relevant_text  # Return error message if retrieve_relevant_text failed
+        input_text = f"question: {question} context: {relevant_text}"
+        response = qa_pipeline(question=question, context=relevant_text)
+        return response['answer']
+    except Exception as e:
+        return str(e)
 
 # Define Gradio interface
 iface = gr.Interface(
     fn=answer_question,
     inputs=[gr.inputs.File(type="file", label="Upload PDF"), gr.inputs.Textbox(lines=2, placeholder="Ask a question")],
     outputs=gr.outputs.Textbox(label="Answer"),
-    title="PDF Q&A with T5"
+    title="PDF Q&A with Mixed Models"
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()