Update app.py

app.py

@@ -1,25 +1,59 @@
+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: 890px !important}
+.gradio-container{max-width: 1000px !important}
 h1{text-align:center}
 footer {
     visibility: hidden
 }
 '''
 
-ACCESS_TOKEN = os.getenv("HF_TOKEN")
+client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
+
+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
+
+# Generate text
+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)
 
-client = OpenAI(
-    base_url="https://api-inference.huggingface.co/v1/",
-    api_key=ACCESS_TOKEN,
-)
+    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)
 
-# Function to save generated text to a file
+    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
+
+# Save the generated content to a file
 def save_file(content, filename, file_format):
     if file_format == "pdf":
         pdf = FPDF()
@@ -42,74 +76,33 @@ def save_file(content, filename, file_format):
     else:
         raise ValueError("Unsupported file format")
 
-# Respond function with file saving
-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-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
-    
-    return response, history + [(message, response)]
-
-# Function to handle file saving after generation
-def save_generated_file(response, filename, file_format):
-    saved_file = save_file(response, filename, file_format)
-    return saved_file
-
-# Gradio interface using Blocks
-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")
-    filename = gr.Textbox(value="output", label="Filename")
-    file_format = gr.Radio(["pdf", "docx", "txt"], label="File Format", value="pdf")
-
-    message = gr.Textbox(label="User Message")
-    chat_history = gr.State(value=[])
-    response_output = gr.Textbox(label="Generated Response")
-    file_output = gr.File(label="Download File")
-
-    generate_button = gr.Button("Generate")
-    save_button = gr.Button("Save to File")
-
-    generate_button.click(
-        respond,
-        inputs=[message, chat_history, system_message, max_tokens, temperature, top_p],
-        outputs=[response_output, chat_history]
-    )
-
-    save_button.click(
-        save_generated_file,
-        inputs=[response_output, filename, file_format],
-        outputs=[file_output]
-    )
+# Combine generate and save file functions
+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
+    # Ensure </s> tags are removed from the final 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
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=generate_and_save,
+    inputs=[
+        gr.Textbox(placeholder="Type your message here...", label="Prompt"),
+        gr.State(value=[]),  # history
+        gr.Textbox(placeholder="Filename (default: output)", label="Filename", value="output"),
+        gr.Radio(["pdf", "docx", "txt"], label="File Format", value="pdf"),
+    ],
+    outputs=[
+        gr.Textbox(label="Generated Text"),
+        gr.State(value=[]),  # history
+        gr.File(label="Download File")
+    ],
+    css=css,
+    title="",
+    theme="bethecloud/storj_theme"
+)
 
-if __name__ == "__main__":
-    demo.launch()
+demo.queue().launch(show_api=False)