Update app.py

app.py

@@ -1,59 +1,62 @@
-import os
-import sys
-import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from PIL import Image
-
-sys.path.append(".")
-os.environ["HUGGINGFACE_HTTPS_PROXY"] = "" # Disable HTTPS proxy when not required
-
-tokenizer = AutoTokenizer.from_pretrained("./model_dir")
-model = AutoModelForSeq2SeqLM.from_pretrained("./model_dir", device_map="auto").half()
-device = next(iter(model.parameters())).device
-
-def generate_response(input_text):
-    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
-    outputs = model.generate(
-        input_ids,
-        max_length=512,
-        num_beams=5,
-        early_stopping=True,
-        pad_token_id=tokenizer.pad_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-        length_penalty=1.0,
-        no_repeat_ngram_size=2,
-        min_length=10,
-        temperature=0.9,
+import streamlit as st
+from huggingface_hub import InferenceClient
+
+client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+
+def format_prompt(message, history):
+    prompt = "<s>"
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
+
+def generate(prompt, history, system_prompt, temperature=0.9, max_new_tokens=256, 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,
     )
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return generated_text
 
-import streamlit as st
+    formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+    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
+    return output
+
+# Create text input for user message
+message_input = st.text_input("You:", "")
+
+# Create text input for system prompt
+system_prompt_input = st.text_input("System Prompt:", "You are a helpful assistant.")
+
+# Create sliders for temperature, max new tokens, top-p, and repetition penalty
+temperature_slider = st.slider("Temperature", 0.0, 1.0, 0.9)
+max_new_tokens_slider = st.slider("Max new tokens", 0, 1048, 256)
+top_p_slider = st.slider("Top-p (nucleus sampling)", 0.0, 1.0, 0.95)
+repetition_penalty_slider = st.slider("Repetition penalty", 1.0, 2.0, 1.0)
+
+# Create button to generate response
+if st.button("Generate"):
+    # Create empty list to store conversation history
+    history = []
+
+    # Call generate function with user message, system prompt, and slider values
+    output = generate(message_input, history, system_prompt_input, temperature=temperature_slider, max_new_tokens=max_new_tokens_slider, top_p=top_p_slider, repetition_penalty=repetition_penalty_slider)
+
+    # Display generated response
+    st.write("Assistant:", output)
 
-st.set_page_config(layout="wide")
-col1, col2 = st.beta_columns((3, 1))
-
-with open('style.css') as f:
-    st.markdown(f'<style>{f.read()}<\style>', unsafe_allow_html=True)
-
-with col1:
-    user_prompt = st.text_area("You:", "", height=50)
-    
-if 'generated' not in st.session_state:
-    st.session_state['generated'] = ''
-
-if len(user_prompt) > 0 and st.button('Send'):
-    response = generate_response(user_prompt)
-    st.write('<span style="font-weight:bold;">Assistant:</span>
-' + response, unsafe_allow_html=True)
-    st.session_state['generated'] += '\n\n<span style="font-weight:bold;">User:</span>'+'\n'+ user_prompt + '\n'
-    st.session_state['generated'] += '<span style="font-weight:bold;">Assistant:</span>\n' + response
-
-if 'generated' in st.session_state:
-    message = st.session_state['generated'].replace('\n', '
-').replace('<span style="font-weight: bold;">User:</span>', '&uarr;').replace('<span style="font-weight: bold;">Assistant:</span>', '')
-    st.markdown(message, unsafe_allow_html=True)
-
-# Upload logo
-logo = Image.open("your_logo.png")
-st.sidebar.image(logo, width=160)
+    # Add user message and generated response to conversation history
+    history.append((message_input, output))