Rehan3024 commited on
Commit
5324b1e
1 Parent(s): a7af8b0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
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.")