Update app.py

app.py

@@ -1,90 +1,69 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
-import torch
+from llama_cpp import Llama
+import os
 
-# App Configuration
-st.set_page_config(page_title="Klove AI ChatBox", page_icon="⚡")
+# Page configuration
+st.set_page_config(page_title="Qwen3 GGUF Chat", page_icon="🤖")
+st.title("🤖 Qwen3-1.7B (8-bit GGUF)")
 
-# 1. Load Qwen3 (Cached for efficiency)
+# 1. Load Model (Cached to prevent reloading on every interaction)
 @st.cache_resource
-def load_qwen3():
-    model_id = "Qwen/Qwen3-1.7B"
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype="auto",
-        device_map="auto"
+def load_qwen_gguf():
+    # repo_id and filename as per your request
+    repo_id = "Qwen/Qwen3-1.7B-GGUF"
+    filename = "Qwen3-1.7B-Q8_0.gguf"
+    
+    # from_pretrained automatically downloads and caches the .gguf file
+    return Llama.from_pretrained(
+        repo_id=repo_id,
+        filename=filename,
+        n_ctx=4096,           # Context window (Qwen3 supports up to 32k, but 4k is safer for free RAM)
+        n_threads=2,          # Match the 2 vCPUs on HF Free Tier
+        verbose=False         # Reduces log clutter
     )
-    return tokenizer, model
-
-tokenizer, model = load_qwen3()
-
-# 2. System Prompt Selection (Main UI, No Sidebar)
-st.title("⚡ Klove AI ChatBox")
-
-system_options = {
-    "General Assistant": "You are a helpful and concise assistant.",
-    "Python Expert": "You are an expert Python developer. Provide clean, efficient code with brief explanations.",
-    "Creative Storyteller": "You are a creative writer. Use vivid imagery and engaging narrative styles."
-}
-
-# Horizontal layout for the selector and a clear button
-col1, col2 = st.columns([3, 1])
-with col1:
-    selected_role = st.selectbox("Choose AI Personality:", list(system_options.keys()))
-with col2:
-    if st.button("Clear History", use_container_width=True):
-        st.session_state.messages = []
-        st.rerun()
 
-system_prompt = system_options[selected_role]
+llm = load_qwen_gguf()
 
-# 3. Setup Session State
+# 2. Chat History Initialization
 if "messages" not in st.session_state:
-    st.session_state.messages = []
+    st.session_state.messages = [
+        {"role": "system", "content": "You are a helpful AI assistant. Answer concisely and accurately."}
+    ]
 
-# Display Chat History
+# Display existing chat messages
 for message in st.session_state.messages:
-    if message["role"] != "system": # Hide system prompt from UI
+    if message["role"] != "system":
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
-# 4. Chat Input & Streaming
-if prompt := st.chat_input("Message Qwen3..."):
-    # Insert system prompt if history is empty
-    if len(st.session_state.messages) == 0:
-        st.session_state.messages.append({"role": "system", "content": system_prompt})
-    
-    # Add User Message
+# 3. Chat Logic & Streaming
+if prompt := st.chat_input("Ask Qwen3..."):
+    # Add user message to history
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
 
-    # Generate Assistant Response
+    # Generate assistant response
     with st.chat_message("assistant"):
-        # Prepare input with chat template
-        input_ids = tokenizer.apply_chat_template(
-            st.session_state.messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors="pt"
-        ).to(model.device)
-
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-        
-        generation_kwargs = dict(
-            input_ids=input_ids,
-            streamer=streamer,
-            max_new_tokens=1024,
-            temperature=0.7
+        # Use create_chat_completion for OpenAI-style streaming
+        stream = llm.create_chat_completion(
+            messages=st.session_state.messages,
+            stream=True,
+            max_tokens=1024,
+            temperature=0.7,
+            top_p=0.8,
+            repeat_penalty=1.1
         )
+        
+        # Generator for Streamlit streaming
+        def stream_response():
+            for chunk in stream:
+                delta = chunk['choices'][0]['delta']
+                if 'content' in delta:
+                    yield delta['content']
 
-        # Threaded generation for real-time streaming
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
-        thread.start()
-
-        # Stream the output
-        full_response = st.write_stream(streamer)
+        # write_stream handles the "typewriter" effect automatically
+        full_response = st.write_stream(stream_response())
     
+    # Save the assistant's final answer
     st.session_state.messages.append({"role": "assistant", "content": full_response})