Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -3,6 +3,7 @@ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModelForQuest
|
|
3 |
from sentence_transformers import SentenceTransformer
|
4 |
import fitz # PyMuPDF
|
5 |
import os
|
|
|
6 |
|
7 |
# Load the models
|
8 |
summarization_model_name = 'facebook/bart-large-cnn'
|
@@ -29,13 +30,13 @@ def summarize_document(document):
|
|
29 |
|
30 |
# Function to get answer to question
|
31 |
def get_answer(question, context):
|
32 |
-
inputs = qa_tokenizer(question, context, return_tensors="pt")
|
33 |
start_positions, end_positions = qa_model(**inputs)
|
34 |
answer_start = torch.argmax(start_positions)
|
35 |
answer_end = torch.argmax(end_positions) + 1
|
36 |
answer = qa_tokenizer.convert_tokens_to_string(qa_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
|
37 |
return answer
|
38 |
-
|
39 |
# Streamlit app
|
40 |
st.title("PDF Summarizer and Q&A")
|
41 |
st.write("Upload a PDF file to get a summary and ask questions about the content.")
|
|
|
3 |
from sentence_transformers import SentenceTransformer
|
4 |
import fitz # PyMuPDF
|
5 |
import os
|
6 |
+
import torch
|
7 |
|
8 |
# Load the models
|
9 |
summarization_model_name = 'facebook/bart-large-cnn'
|
|
|
30 |
|
31 |
# Function to get answer to question
|
32 |
def get_answer(question, context):
|
33 |
+
inputs = qa_tokenizer(question, context, return_tensors="pt", padding=True, truncation=True, max_length=512)
|
34 |
start_positions, end_positions = qa_model(**inputs)
|
35 |
answer_start = torch.argmax(start_positions)
|
36 |
answer_end = torch.argmax(end_positions) + 1
|
37 |
answer = qa_tokenizer.convert_tokens_to_string(qa_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
|
38 |
return answer
|
39 |
+
|
40 |
# Streamlit app
|
41 |
st.title("PDF Summarizer and Q&A")
|
42 |
st.write("Upload a PDF file to get a summary and ask questions about the content.")
|