Spaces:
Sleeping
Sleeping
File size: 2,948 Bytes
10d7745 43c74e3 10d7745 43c74e3 10d7745 df61c4d 43c74e3 10d7745 43c74e3 df61c4d 43c74e3 d39f2ea 43c74e3 10d7745 43c74e3 df61c4d 43c74e3 df61c4d 43c74e3 d39f2ea 43c74e3 d39f2ea df61c4d 43c74e3 d39f2ea 43c74e3 d39f2ea df61c4d 43c74e3 10d7745 df61c4d 43c74e3 df61c4d 43c74e3 df61c4d 43c74e3 df61c4d 43c74e3 df61c4d 43c74e3 df61c4d 43c74e3 df61c4d d39f2ea cfdd22d 43c74e3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import fitz # PyMuPDF for PDF handling
# Load the model and tokenizer
@st.cache_resource
def load_model():
# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("himmeow/vi-gemma-2b-RAG")
model = AutoModelForCausalLM.from_pretrained(
"himmeow/vi-gemma-2b-RAG",
device_map="auto",
torch_dtype=torch.bfloat16
)
if torch.cuda.is_available():
model.to("cuda")
return tokenizer, model
# Function to extract text from PDF
def extract_text_from_pdf(pdf_file):
# Extract text from the uploaded PDF file using PyMuPDF
doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
text = ""
for page_num in range(doc.page_count):
page = doc.load_page(page_num)
text += page.get_text("text") # Ensure text extraction
return text
# Function to generate response from model
def generate_response(input_text, query, tokenizer, model):
# Format the input prompt for the model
prompt = f"""
### Instruction and Input:
Based on the following context/document:
{input_text}
Please answer the question: {query}
### Response:
"""
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
if torch.cuda.is_available():
input_ids = input_ids.to("cuda")
# Generate a response from the model
outputs = model.generate(
input_ids=input_ids,
max_new_tokens=500,
no_repeat_ngram_size=5
)
# Decode the generated output into readable text
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Streamlit app main function
def main():
st.title("PDF Question Answering with vi-gemma-2b-RAG")
# File uploader widget for PDF files
pdf_file = st.file_uploader("Upload a PDF file", type=["pdf"])
if pdf_file is not None:
with st.spinner("Reading the PDF..."):
# Extract text from the uploaded PDF
pdf_text = extract_text_from_pdf(pdf_file)
st.text_area("Extracted Text", pdf_text, height=300)
# Text input for the user's question
query = st.text_input("Enter your question:")
if st.button("Get Answer"):
if query.strip() == "":
st.warning("Please enter a question.")
else:
with st.spinner("Generating response..."):
# Load the model and tokenizer
tokenizer, model = load_model()
# Generate the response using the model
try:
response = generate_response(pdf_text, query, tokenizer, model)
st.text_area("Response", response, height=200)
except Exception as e:
st.error(f"Error generating response: {e}")
if __name__ == "__main__":
main()
|