Update app.py

app.py

@@ -1,58 +1,25 @@
-from huggingface_hub import InferenceClient
 import gradio as gr
+from openai import OpenAI
+import os
 from fpdf import FPDF
 import docx
 
 css = '''
-.gradio-container{max-width: 1000px !important}
+.gradio-container{max-width: 890px !important}
 h1{text-align:center}
 footer {
     visibility: hidden
 }
 '''
 
-client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
+ACCESS_TOKEN = os.getenv("HF_TOKEN")
 
-def format_prompt(message, history, system_prompt=None):
-    prompt = "<s>"
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    if system_prompt:
-        prompt += f"[SYS] {system_prompt} [/SYS]"
-    prompt += f"[INST] {message} [/INST]"
-    return prompt
-
-def generate(
-    prompt, history, system_prompt=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0,
-):
-    temperature = float(temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(top_p)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42,
-    )
-
-    formatted_prompt = format_prompt(prompt, history, system_prompt)
-
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-        # Clean up </s> tags from the generated output
-        output = output.replace("</s>", "")
-        yield output
-    return output
+client = OpenAI(
+    base_url="https://api-inference.huggingface.co/v1/",
+    api_key=ACCESS_TOKEN,
+)
 
-# Save the generated content to a file
+# Function to save generated text to a file
 def save_file(content, filename, file_format):
     if file_format == "pdf":
         pdf = FPDF()
@@ -75,30 +42,70 @@ def save_file(content, filename, file_format):
     else:
         raise ValueError("Unsupported file format")
 
-def generate_and_save(prompt, history, filename="output", file_format="pdf", system_prompt=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0):
-    generated_text = ""
-    for output in generate(prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty):
-        generated_text = output
-    generated_text = generated_text.replace("</s>", "")
-    saved_file = save_file(generated_text, filename, file_format)
-    return generated_text, history + [(prompt, generated_text)], saved_file
+# Respond function with file saving
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    filename,
+    file_format
+):
+    messages = [{"role": "system", "content": system_message}]
+
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
+
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+
+    for message in client.chat.completions.create(
+        model="meta-llama/Meta-Llama-3.1-70B-Instruct",
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+        messages=messages,
+    ):
+        token = message.choices[0].delta.content
+        response += token
+        yield response
     
-demo = gr.Interface(
-    fn=generate_and_save,
-    inputs=[
-        gr.Textbox(placeholder="Type your message here...", label="Chatbot", lines=4),
-        gr.State(value=[]), 
-        gr.Textbox(placeholder="Filename (default: output)", label="Filename", value="output"),
+    # Save the final response to the specified file format
+    saved_file = save_file(response, filename, file_format)
+    yield response, history + [(message, response)], saved_file
+
+# Gradio interface
+demo = gr.ChatInterface(
+    fn=respond,
+    additional_inputs=[
+        gr.Textbox(value="", 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",
+        ),
+        gr.Textbox(value="output", label="Filename"),
         gr.Radio(["pdf", "docx", "txt"], label="File Format", value="pdf"),
     ],
     outputs=[
-        gr.Textbox(label="Generated Text", lines=4),
-        gr.State(value=[]), 
-        gr.File(label="Download File")
+        gr.Textbox(label="Generated Text"),
+        gr.State(value=[]),  # history
+        gr.File(label="Download File"),
     ],
     css=css,
-    title="GRAB DOC",
-    theme="bethecloud/storj_theme"
+    theme="allenai/gradio-theme",
 )
 
-demo.queue().launch(show_api=False)
+if __name__ == "__main__":
+    demo.launch()