Update app.py

app.py

@@ -1,32 +1,39 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import bitsandbytes as bnb
-import torch
+from huggingface_hub import InferenceClient
+import os
 
-# Load the model and tokenizer with 4-bit quantization
-@st.cache_resource
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-7B", trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        "Qwen/Qwen-7B",
-        load_in_4bit=True,
-        device_map="auto",
-        trust_remote_code=True  # Allows custom code execution
-    )
-    return tokenizer, model
+# Retrieve the Hugging Face token from environment variables
+hf_token = os.getenv("HF_API_TOKEN")
 
-tokenizer, model = load_model()
+# Initialize the Hugging Face Inference Client
+client = InferenceClient(
+    model="Qwen/Qwen2-7B",
+    token=hf_token,
+)
 
-# Streamlit app UI
-st.title("Qwen-7B Text Generation with 4-bit Quantization")
+def get_chat_completion(message):
+    try:
+        response = client.chat_completion(
+            messages=[{"role": "user", "content": message}],
+            max_tokens=500,
+            stream=False,
+        )
+        completion = response[0]['choices'][0]['message']['content']
+        return completion
+    except Exception as e:
+        return f"Error: {e}"
 
-# Text input
-user_input = st.text_area("Enter your text:")
+# Streamlit app layout
+st.title("Chat with Hugging Face Model")
 
-# Generate text on button click
-if st.button("Generate"):
-    inputs = tokenizer(user_input, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=100)
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    st.write("Generated Text:")
-    st.write(generated_text)
+# Input from the user
+user_input = st.text_input("Enter your message:")
+
+if st.button("Send"):
+    if user_input:
+        # Get response from the model
+        response = get_chat_completion(user_input)
+        st.write("**Response:**")
+        st.write(response)
+    else:
+        st.write("Please enter a message.")