app.py

app.py

@@ -0,0 +1,49 @@
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load the model and tokenizer
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
+    model = AutoModelForCausalLM.from_pretrained(
+        "meta-llama/Llama-2-7b-chat-hf", 
+        device_map="auto", 
+        torch_dtype="auto"
+    )
+    return model, tokenizer
+
+# Initialize
+st.title("LLaMA-2 Chatbot")
+st.sidebar.title("Configuration")
+
+# Input parameters
+max_tokens = st.sidebar.slider("Max Tokens", 50, 1000, 256)
+temperature = st.sidebar.slider("Temperature", 0.1, 1.5, 0.7)
+model, tokenizer = load_model()
+
+# Chat interface
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+user_input = st.text_input("Your message:", key="input")
+if st.button("Send"):
+    if user_input:
+        st.session_state.messages.append({"role": "user", "content": user_input})
+
+        inputs = tokenizer(
+            user_input, return_tensors="pt", truncation=True
+        ).to(model.device)
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=max_tokens,
+            temperature=temperature,
+            do_sample=True,
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
+# Display the conversation
+for message in st.session_state.messages:
+    role = "User" if message["role"] == "user" else "Assistant"
+    st.write(f"**{role}:** {message['content']}")