Update app.py

app.py

@@ -1,6 +1,5 @@
+from huggingface_hub import InferenceClient
 import gradio as gr
-from openai import OpenAI
-import os
 from fpdf import FPDF
 from docx import Document
 
@@ -12,106 +11,85 @@ footer {
 }
 '''
 
-ACCESS_TOKEN = os.getenv("HF_TOKEN")
+client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
 
-client = OpenAI(
-    base_url="https://api-inference.huggingface.co/v1/",
-    api_key=ACCESS_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 respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    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-8B-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
-
-def save_to_file(history, file_format):
-    if file_format == "PDF":
+def save_to_file(content, filename, format):
+    if format == "pdf":
         pdf = FPDF()
         pdf.add_page()
         pdf.set_auto_page_break(auto=True, margin=15)
         pdf.set_font("Arial", size=12)
-        for user_message, assistant_message in history:
-            pdf.multi_cell(0, 10, f"User: {user_message}")
-            pdf.multi_cell(0, 10, f"Assistant: {assistant_message}")
-        file_name = "chat_history.pdf"
-        pdf.output(file_name)
-        
-    elif file_format == "DOCX":
+        pdf.multi_cell(0, 10, content)
+        pdf.output(filename)
+    elif format == "docx":
         doc = Document()
-        for user_message, assistant_message in history:
-            doc.add_paragraph(f"User: {user_message}")
-            doc.add_paragraph(f"Assistant: {assistant_message}")
-        file_name = "chat_history.docx"
-        doc.save(file_name)
-        
-    elif file_format == "TXT":
-        file_name = "chat_history.txt"
-        with open(file_name, "w") as file:
-            for user_message, assistant_message in history:
-                file.write(f"User: {user_message}\n")
-                file.write(f"Assistant: {assistant_message}\n")
-                
-    return file_name
+        doc.add_paragraph(content)
+        doc.save(filename)
+    elif format == "txt":
+        with open(filename, 'w') as file:
+            file.write(content)
 
-# Gradio Interface Setup
-with gr.Blocks(css=css) as demo:
-    system_message = gr.Textbox(value="", label="System message")
-    max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
-    temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
-    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
-    save_as = gr.Radio(["PDF", "DOCX", "TXT"], label="Save As")
-    
-    chat = gr.Chatbot()
-    msg = gr.Textbox(label="Your message")
+def generate(
+    prompt, history, system_prompt=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0,
+    save_format="txt", save_file=None
+):
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
 
-    def respond_wrapper(message, history):
-        response_generator = respond(
-            message,
-            history,
-            system_message.value,
-            max_tokens.value,
-            temperature.value,
-            top_p.value
-        )
-        response = next(response_generator)
-        return history + [(message, response)]
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=42,
+    )
 
-    msg.submit(respond_wrapper, [msg, chat], [chat])
-    
-    save_button = gr.Button("Save Conversation")
-    output_file = gr.File(label="Download File")
+    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
+        yield output
     
-    def handle_save(history, file_format):
-        return save_to_file(history, file_format)
+    # Save content to file
+    if save_file:
+        save_to_file(output, save_file, save_format)
+    return output
 
-    save_button.click(handle_save, inputs=[chat, save_as], outputs=output_file)
+with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+    gr.Markdown("# GRAB DOC")
+    prompt = gr.Textbox(label="Prompt", placeholder="Enter your text here...", lines=4)
+    history = gr.State([])
+    system_prompt = gr.Textbox(label="System Prompt", placeholder="Enter system instructions (optional)...", lines=2, visible=False)
+    temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, step=0.1, value=0.2)
+    max_new_tokens = gr.Slider(label="Max New Tokens", minimum=1, maximum=1024, step=1, value=1024)
+    top_p = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, step=0.05, value=0.95)
+    repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=0.0, maximum=2.0, step=0.1, value=1.0)
+    save_format = gr.Dropdown(label="Save Format", choices=["txt", "pdf", "docx"], value="txt")
+    save_file = gr.Textbox(label="Save File Name", placeholder="Enter the filename (without extension)...", lines=1)
+    
+    submit = gr.Button("Generate")
+    output = gr.Textbox(label="Generated Output", lines=20)
+    
+    submit.click(
+        generate, 
+        inputs=[prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty, save_format, save_file], 
+        outputs=output
+    )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.queue().launch(show_api=False)