Update app.py

app.py

@@ -1,4 +1,69 @@
 import streamlit as st
+import time
+import requests
+import os
+from huggingface_hub import InferenceClient
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+# Hugging Face API Setup
+API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN")
+GPT2XL_API_URL = "https://api-inference.huggingface.co/models/openai-community/gpt2-xl"
+MISTRAL_MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
+client = InferenceClient(api_key=API_TOKEN)
+
+# Query GPT-2 XL
+def query_from_gpt2xl(text: str):
+    headers = {"Authorization": f"Bearer {API_TOKEN}"}
+    while True:
+        response = requests.post(GPT2XL_API_URL, headers=headers, json={"inputs": text})
+        response_data = response.json()
+        if "error" in response_data and "loading" in response_data["error"]:
+            wait_time = response_data.get("estimated_time", 10)
+            st.info(f"Model is loading. Waiting for {wait_time:.2f} seconds...")
+            time.sleep(wait_time)
+        else:
+            return response_data[0]["generated_text"]
+
+# Query Mistral
+def query_from_mistral(text: str):
+    messages = [{"role": "user", "content": text}]
+    completion = client.chat.completions.create(
+        model=MISTRAL_MODEL_NAME,
+        messages=messages,
+        max_tokens=500,
+    )
+    return completion.choices[0].message["content"]
+
+def main():
+    st.set_page_config(page_title="Multi-Model Chat", layout="centered")
+    st.title("🤖 Multi-Model Chat")
+    st.markdown("Chat with either **GPT-2 XL** or **Mistral-7B-Instruct** via Hugging Face API.")
+
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    model_choice = st.selectbox("Select a model:", ["GPT-2 XL", "Mistral-7B-Instruct"])
+
+    with st.form(key="chat_form", clear_on_submit=True):
+        user_input = st.text_input("You:", "")
+        submit = st.form_submit_button("Send")
+
+    if submit and user_input:
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        with st.spinner("Fetching response..."):
+            try:
+                if model_choice == "GPT-2 XL":
+                    response = query_from_gpt2xl(user_input)
+                elif model_choice == "Mistral-7B-Instruct":
+                    response = query_from_mistral(user_input)
+                st.session_state.messages.append({"role": "bot", "content": response})
+            except Exception as e:
+                st.error(f"Error: {e}")
+
+    for message in st.session_state.messages:
+        if message["role"] == "user":
+            st.markdown(f"**You:** {message['content']}")
+        elif message["role"] == "bot":
+            st.markdown(f"**Bot:** {message['content']}")
+
+if __name__ == "__main__":
+    main()