Spaces:
Runtime error
Runtime error
import gradio as gr | |
import PyPDF2 | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
import torch | |
# Initialize the model and tokenizer | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
model = AutoModelForCausalLM.from_pretrained( | |
"pints-ai/1.5-Pints-2k-v0.1", | |
device_map=device, | |
# attn_implementation="flash_attention_2" | |
) | |
tokenizer = AutoTokenizer.from_pretrained("pints-ai/1.5-Pints-2k-v0.1") | |
def extract_pdf_content(file): | |
pdf_reader = PyPDF2.PdfReader(file) | |
content = "" | |
for page in pdf_reader.pages: | |
content += page.extract_text() + "\n" | |
return content | |
def chat(message, history, pdf_content): | |
# Construct the full conversation history | |
full_history = [ | |
{"role": "system", "content": f"You are an AI assistant that follows instruction extremely well. Help as much as you can. Use the following information from the uploaded PDF as context: {pdf_content}"} | |
] + history + [{"role": "user", "content": message}] | |
text = tokenizer.apply_chat_template( | |
full_history, | |
tokenize=False, | |
add_generation_prompt=True | |
) | |
input = tokenizer([text], return_tensors="pt").to(device) | |
generated_ids = model.generate( | |
input.input_ids, | |
max_new_tokens=512 | |
) | |
input_length = len(input.input_ids[0]) | |
response = tokenizer.decode(generated_ids[0][input_length:]) | |
# Update the history with the new message and response | |
history.append({"role": "user", "content": message}) | |
history.append({"role": "assistant", "content": response}) | |
return history, [(history[i]["content"], history[i+1]["content"]) for i in range(0, len(history)-1, 2)] | |
def process_pdf(file): | |
if file is None: | |
return "Please upload a PDF file.", "" | |
content = extract_pdf_content(file) | |
return f"PDF content extracted. Length: {len(content)} characters.", content | |
with gr.Blocks() as demo: | |
pdf_content = gr.State("") | |
chat_history = gr.State([]) | |
with gr.Row(): | |
pdf_upload = gr.File(label="Upload PDF", file_types=[".pdf"]) | |
pdf_output = gr.Textbox(label="PDF Processing Output") | |
chatbot = gr.Chatbot() | |
msg = gr.Textbox() | |
pdf_upload.upload(process_pdf, pdf_upload, [pdf_output, pdf_content]) | |
msg.submit(chat, [msg, chat_history, pdf_content], [chat_history, chatbot]) | |
demo.launch() |