Update app.py

app.py

@@ -1,6 +1,6 @@
 from huggingface_hub import InferenceClient
 import gradio as gr
-import re
+import PyPDF2
 
 client = InferenceClient(
     "mistralai/Mixtral-8x7B-Instruct-v0.1"
@@ -14,7 +14,9 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def generate(prompt, history, system_prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, file=None):
+def generate(
+    prompt, history, system_prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, file=None
+):
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
@@ -30,18 +32,24 @@ def generate(prompt, history, system_prompt, temperature=0.9, max_new_tokens=256
     )
 
     if file:
-        sentences = process_file(file)
-        prompt = "\n".join(sentences)
+        text = extract_text_from_pdf(file)
+        prompt = text
 
     formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
-    response = client.text_generation(formatted_prompt, **generate_kwargs, details=True, return_full_text=True)
-    return response.text
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    output = ""
 
-def process_file(file):
-    text = file.decode("utf-8")
-    sentences = re.split(r'[.!?]+', text)
-    sentences = [s.strip() for s in sentences if s.strip()]
-    return sentences
+    for response in stream:
+        output += response.token.text
+        yield output
+    return output
+
+def extract_text_from_pdf(file):
+    pdf_reader = PyPDF2.PdfReader(file)
+    text = ""
+    for page in range(len(pdf_reader.pages)):
+        text += pdf_reader.pages[page].extract_text()
+    return text
 
 additional_inputs=[
     gr.Textbox(
@@ -85,7 +93,7 @@ additional_inputs=[
         interactive=True,
         info="Penalize repeated tokens",
     ),
-    gr.File(label="Upload File", file_count="single"),
+    gr.File(label="Upload PDF File", file_count="single", file_types=[".pdf"]),
 ]
 
 gr.ChatInterface(