File size: 4,867 Bytes
7aa1c24
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f1bb9f3
 
7aa1c24
 
 
 
 
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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
# app.py
import gradio as gr
from huggingface_hub import InferenceClient
import PyPDF2
import io

# How many characters to include per uploaded PDF (avoid huge inputs / token blowup)
MAX_CHARS_PER_PDF = 20000

def extract_text_from_pdf_fileobj(file_obj) -> str:
    """
    Extract text from a file-like object containing a PDF (Gradio gives a tempfile path
    accessible via file.name, but file might also be an in-memory BytesIO).
    Returns the first MAX_CHARS_PER_PDF characters.
    """
    try:
        # file_obj may be a _TemporaryFileWrapper with .name path or a dict/file-like; handle both
        if hasattr(file_obj, "name"):
            reader = PyPDF2.PdfReader(file_obj.name)
        else:
            # try to read bytes
            file_obj.seek(0)
            reader = PyPDF2.PdfReader(io.BytesIO(file_obj.read()))
    except Exception as e:
        # best-effort: return empty string if cannot parse
        return f"[Could not read PDF: {e}]"

    text_parts = []
    try:
        for page in reader.pages:
            try:
                text_parts.append(page.extract_text() or "")
            except Exception:
                # ignore pages that fail
                continue
    except Exception:
        # some PDFs may throw during iteration — fallback
        pass

    combined = "\n".join(text_parts)
    if len(combined) > MAX_CHARS_PER_PDF:
        combined = combined[:MAX_CHARS_PER_PDF] + "\n\n[TRUNCATED]"
    return combined

def build_context_from_uploaded_files(uploaded_files):
    """
    Given the list returned by gr.File (list of file-like objects), return a single string
    that summarizes / contains the extracted text to be sent as extra context.
    """
    if not uploaded_files:
        return ""

    ctx_parts = []
    for f in uploaded_files:
        try:
            # `f` is a TemporaryFile object from gradio with .name attribute
            extracted = extract_text_from_pdf_fileobj(f)
            header = f"--- Begin extracted text from uploaded file: {getattr(f, 'name', 'uploaded_pdf')} ---\n"
            footer = f"\n--- End of {getattr(f, 'name', 'uploaded_pdf')} ---\n\n"
            ctx_parts.append(header + extracted + footer)
        except Exception as e:
            ctx_parts.append(f"[Error extracting {getattr(f, 'name', 'uploaded_pdf')}: {e}]")

    return "\n".join(ctx_parts)


def respond(
    message,
    history: list[dict[str, str]],
    uploaded_files,                 # NEW: list of uploaded PDF files
    system_message,
    max_tokens,
    temperature,
    top_p,
    hf_token: gr.OAuthToken,
):
    """
    For more information on `huggingface_hub` Inference API support, please check the docs:
    https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
    """
    # Build an extra 'context' system message from uploaded PDFs (if any)
    uploaded_context = build_context_from_uploaded_files(uploaded_files)
    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")

    # system message first
    messages = [{"role": "system", "content": system_message}]

    # if we have uploaded PDF content, add it as another system message so the model sees it
    if uploaded_context:
        messages.append({"role": "system", "content": "Context from uploaded PDFs:\n\n" + uploaded_context})

    # replay conversation history (assumed to be a list of role/content dicts)
    messages.extend(history)

    # finally the user message
    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        choices = message.choices
        token = ""
        if len(choices) and choices[0].delta.content:
            token = choices[0].delta.content

        response += token
        yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[
        # NEW: add a File uploader as the first additional input
        gr.File(label="Upload PDFs (optional)", file_count="multiple", file_types=[".pdf"]),
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)

with gr.Blocks() as demo:
    with gr.Sidebar():
        gr.LoginButton()
    chatbot.render()


if __name__ == "__main__":
    demo.launch()