Update app.py

app.py

@@ -1,50 +1,57 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import snapshot_download, login
-from pathlib import Path
-
-def main():
-    st.title("Codestral Inference with Hugging Face")
-    
-    # Get the Hugging Face API token from the user
-    hf_token = st.text_input("Enter your Hugging Face API token", type="password")
-    if not hf_token:
-        st.warning("Please enter your Hugging Face API token to proceed.")
-        st.stop()
-
-    # Login to Hugging Face Hub
-    login(hf_token)
-    
-    # Download the model files
-    st.text("Downloading model...")
-    model_id = "mistralai/Codestral-22B-v0.1"
-    local_model_path = Path.home().joinpath('mistral_models', model_id)
-    local_model_path.mkdir(parents=True, exist_ok=True)
-
-    snapshot_download(repo_id=model_id, local_dir=local_model_path, use_auth_token=hf_token)
-    st.success("Model downloaded successfully!")
-    
-    # Load the model and tokenizer
-    st.text("Loading model...")
-    tokenizer = AutoTokenizer.from_pretrained(local_model_path, use_auth_token=hf_token)
-    model = AutoModelForCausalLM.from_pretrained(local_model_path, use_auth_token=hf_token)
-    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    st.success("Model loaded successfully!")
-    
-    user_input = st.text_area("Enter your instruction", "Explain Machine Learning to me in a nutshell.")
-    max_tokens = st.slider("Max Tokens", min_value=10, max_value=500, value=64)
-    temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.7)
-    
-    if st.button("Generate"):
-        with st.spinner("Generating response..."):
-            result = generate_response(generator, user_input, max_tokens, temperature)
-            st.success("Response generated!")
-            st.text_area("Generated Response", result, height=200)
-
-def generate_response(generator, user_input, max_tokens, temperature):
-    response = generator(user_input, max_new_tokens=max_tokens, do_sample=True, temperature=temperature)
-    result = response[0]['generated_text']
-    return result
-
-if __name__ == "__main__":
-    main()
+from huggingface_hub import InferenceClient
+
+"""
+For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+"""
+client = InferenceClient(model_id="mistralai/Codestral-22B-v0.1", token="your_token_here")
+
+def respond(message, history, system_message, max_tokens, temperature, top_p):
+    messages = [{"role": "system", "content": system_message}]
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+    try:
+        for message in client.chat_completion(
+            messages=messages,
+            max_tokens=max_tokens,
+            stream=True,
+            temperature=temperature,
+            top_p=top_p,
+        ):
+            token = message.choices[0].delta.content
+            response += token
+            yield response
+    except Exception as e:
+        yield f"Error: {e}"
+
+# Streamlit interface
+st.title("Chat with Codestral Model")
+system_message = st.text_input("System message", value="You are an expert python coder with in depth knowledge of langchain.")
+max_tokens = st.slider("Max new tokens", min_value=1, max_value=2048, value=2048, step=1)
+temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.6, step=0.1)
+top_p = st.slider("Top-p (nucleus sampling)", min_value=0.1, max_value=1.0, value=0.95, step=0.05)
+
+history = []
+
+if "history" not in st.session_state:
+    st.session_state.history = []
+
+def get_response():
+    user_input = st.session_state.user_input
+    if user_input:
+        st.session_state.history.append((user_input, ""))
+        response_generator = respond(user_input, st.session_state.history, system_message, max_tokens, temperature, top_p)
+        response = ""
+        for r in response_generator:
+            response = r
+        st.session_state.history[-1] = (user_input, response)
+
+st.text_area("Chat History", value="\n".join([f"User: {h[0]}\nAssistant: {h[1]}" for h in st.session_state.history]), height=300)
+
+st.text_input("Your message:", key="user_input", on_change=get_response)