Update app.py

app.py

@@ -1,13 +1,52 @@
 import streamlit as st
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import InferenceClient
+import os
 
-# Load the Mistral model and tokenizer
-model_name = "mistralai/Mistral-7B-Instruct-v0.2"
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.1")
 
-# Initialize the text generation pipeline
-generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
+def format_prompt(message, history, system_prompt=""):
+    prompt = "<s>"
+    if system_prompt:
+        prompt += f"[SYS] {system_prompt} [/SYS] "
+    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 extract_filename_and_code(bot_response):
+    try:
+        with open("default.py", 'w') as file:
+            file.write(bot_response)
+        st.write("Successfully wrote to default.py!")
+    except Exception as e:
+        st.write(f"Error writing to file: {e}")
+
+def generate(prompt, history, system_prompt="", temperature=0.9, max_new_tokens=512, 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
+        extract_filename_and_code(output)
+        yield output
+
+    return output
 
 # Streamlit UI
 st.title("Waifu Character Generator")
@@ -17,13 +56,20 @@ name = st.text_input("Name of the Waifu")
 hair_color = st.selectbox("Hair Color", ["Blonde", "Brunette", "Red", "Black", "Blue", "Pink"])
 personality = st.selectbox("Personality", ["Tsundere", "Yandere", "Kuudere", "Dandere", "Genki", "Normal"])
 outfit_style = st.selectbox("Outfit Style", ["School Uniform", "Maid Outfit", "Casual", "Kimono", "Gothic Lolita"])
+system_prompt = st.text_input("Optional System Prompt", "")
+
+# Advanced settings
+temperature = st.slider("Temperature", 0.0, 1.0, 0.9, step=0.05)
+max_new_tokens = st.slider("Max new tokens", 0, 8192, 512, step=64)
+top_p = st.slider("Top-p (nucleus sampling)", 0.0, 1.0, 0.95, step=0.05)
+repetition_penalty = st.slider("Repetition penalty", 1.0, 2.0, 1.0, step=0.05)
 
 # Generate button
 if st.button("Generate Waifu"):
-    # Generate character description
+    history = []
     prompt = f"Create a waifu character named {name} with {hair_color} hair, a {personality} personality, and wearing a {outfit_style}."
-    result = generator(prompt, max_length=150, num_return_sequences=1)[0]['generated_text']
+    output = generate(prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
     
     # Display the generated character
     st.subheader("Generated Waifu Character")
-    st.write(result)
+    st.write(next(output))