Fix inputs, add OAuth token

app.py

@@ -1,60 +1,83 @@
 import gradio as gr
-from huggingface_hub import InferenceClient, OAuthToken
-
-MODEL_ID = "Bocklitz-Lab/lit2vec-tldr-bart-model"
-
+from huggingface_hub import InferenceClient
+from typing import List, Dict
 
+# Response function for the chatbot
 def respond(
     message: str,
-    history: list[dict[str, str]],
+    history: List[Dict[str, str]],
     system_message: str,
     max_tokens: int,
     temperature: float,
     top_p: float,
-    hf_token: OAuthToken | None,
+    hf_token: gr.OAuthToken,
 ):
-    """Summarise a chemistry abstract with the HF Inference API."""
+    """
+    Sends a chat message to the Hugging Face Inference API using the provided token and parameters.
+    """
     client = InferenceClient(
-        model=MODEL_ID,
-        token=None if hf_token is None else hf_token.token,
+        token=hf_token.token,
+        model="Bocklitz-Lab/lit2vec-tldr-bart-model"
     )
 
-    prompt = f"{system_message.strip()}\n\n{message.strip()}"
+    messages = [{"role": "system", "content": system_message}] + history
+    messages.append({"role": "user", "content": message})
+
+    response = ""
 
-    # stream=False → one-shot; stream=True → token generator
-    for chunk in client.text_generation(
-        prompt,
-        max_new_tokens=max_tokens,
+    for message_chunk in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
         temperature=temperature,
         top_p=top_p,
-        stream=True,            # change to False if you prefer
     ):
-        yield chunk
-
-
-with gr.Blocks(title="🧪 Chemistry Abstract Summariser") as demo:
-    with gr.Sidebar():
-        login_btn = gr.LoginButton()
+        if message_chunk.choices and message_chunk.choices[0].delta.content:
+            token = message_chunk.choices[0].delta.content
+            response += token
+            yield response
 
-    chatbot = gr.ChatInterface(
-        respond,
-        chatbot=gr.Chatbot(type="messages"),
-        textbox=gr.Textbox(
-            placeholder="Paste abstract of a chemistry paper…",
-            lines=8,
+# Define the Gradio interface
+chatbot = gr.ChatInterface(
+    fn=respond,
+    type="messages",
+    additional_inputs=[
+        gr.Textbox(
+            value="You are a friendly chatbot.",
+            label="System message",
+            lines=1
         ),
-        additional_inputs=[
-            gr.Textbox(
-                value="Summarise this chemistry paper abstract:",
-                label="System message",
-            ),
-            gr.Slider(16, 1024, value=256, step=8, label="Max new tokens"),
-            gr.Slider(0.1, 4.0, value=0.7,  step=0.1, label="Temperature"),
-            gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p"),
-            login_btn,                          # 👈 passes OAuthToken to respond()
-        ],
-        type="messages",
-    )
+        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)"
+        ),
+    ],
+)
+
+# Set up the full Gradio Blocks layout with login
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.LoginButton()
+    chatbot.render()
 
+# Run the app
 if __name__ == "__main__":
-    demo.launch()        # Hugging Face Spaces picks this up automatically
+    demo.launch()