Update app.py

app.py

@@ -1,11 +1,17 @@
 import gradio as gr
 import os
-from huggingface_hub import InferenceClient, login
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, pipeline
+from threading import Thread
 import random
 
-# Authenticate using secret environment variable 🔒
-login(token=os.environ.get("HF_TOKEN"))
+# Local model setup 🤖
+model_name = "HuggingFaceH4/zephyr-7b-beta"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    load_in_4bit=True  # Reduces VRAM usage
+)
 
 # Safety tools 🛡️
 BLOCKED_WORDS = ["violence", "hate", "gun", "personal"]
@@ -14,12 +20,8 @@ SAFE_IDEAS = [
     "Code a game about recycling ♻️",
     "Plan an AI tool for school safety 🚸"
 ]
-
-# Safety model
 safety_checker = pipeline("text-classification", model="unitary/toxic-bert")
 
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
 def is_safe(text):
     text = text.lower()
     if any(bad_word in text for bad_word in BLOCKED_WORDS):
@@ -28,16 +30,14 @@ def is_safe(text):
     return not (result["label"] == "toxic" and result["score"] > 0.7)
 
 def respond(message, history, system_message, max_tokens, temperature, top_p):
+    # Safety check first 🔒
     if not is_safe(message):
         return f"🚫 Let's focus on positive projects! Try: {random.choice(SAFE_IDEAS)}"
     
-    messages = [{
-        "role": "system",
-        "content": f"{system_message}\nYou are a friendly STEM mentor for kids. Never discuss unsafe topics!"
-    }]
+    # Prepare chat history
+    messages = [{"role": "system", "content": system_message}]
     
-    # Rest of chat implementation
-    for user_msg, bot_msg in history:
+    for user_msg, bot_msg in history[-5:]:  # Keep last 5 exchanges
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
         if bot_msg:
@@ -45,17 +45,30 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
     
     messages.append({"role": "user", "content": message})
     
-    response = ""
-    for chunk in client.chat_completion(
+    # Tokenize and prepare streaming
+    inputs = tokenizer.apply_chat_template(
         messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p
-    ):
-        token = chunk.choices[0].delta.content
-        response += token
-        yield response
+        return_tensors="pt"
+    ).to(model.device)
+    
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
+    generation_kwargs = {
+        "inputs": inputs,
+        "max_new_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "streamer": streamer
+    }
+    
+    # Start generation in thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    # Stream output
+    partial_message = ""
+    for new_token in streamer:
+        partial_message += new_token
+        yield partial_message
 
 with gr.Blocks() as demo:
     gr.Markdown("# 🤖 REACT Ethical AI Lab")
@@ -74,4 +87,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0")