Update app.py

app.py

@@ -1,6 +1,24 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-#from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
+from threading import Thread
+
+# Load model at startup
+MODEL_NAME = "nosadaniel/llama3-1-8b-tuned"
+print("Loading model and tokenizer...")
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    low_cpu_mem_usage=True
+)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+print("Model loaded successfully!")
 
 def respond(
     message,
@@ -9,43 +27,45 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-    hf_token: gr.OAuthToken,
 ):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    #client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
-    client = InferenceClient(token=hf_token.token, model="nosadaniel/llama3-1-8b-tuned")
-    #model="openai/gpt-oss-20b"
-    #model="nosadaniel/llama3-1-8b-tuned"
-
-    messages = [{"role": "system", "content": system_message}]
-
-    messages.extend(history)
-
-    messages.append({"role": "user", "content": message})
-
+    # Build conversation
+    conversation = f"{system_message}\n\n"
+    
+    for msg in history:
+        if msg["role"] == "user":
+            conversation += f"User: {msg['content']}\n"
+        elif msg["role"] == "assistant":
+            conversation += f"Assistant: {msg['content']}\n"
+    
+    conversation += f"User: {message}\nAssistant:"
+    
+    # Tokenize
+    inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
+    
+    # Setup streamer
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    generation_kwargs = {
+        **inputs,
+        "max_new_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "do_sample": True,
+        "pad_token_id": tokenizer.pad_token_id,
+        "eos_token_id": tokenizer.eos_token_id,
+        "streamer": streamer,
+    }
+    
+    # Generate in separate thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    # Stream output
     response = ""
-
-    for message in client.text_generation(
-        messages,
-        max_new_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
-        response += token
+    for new_text in streamer:
+        response += new_text
         yield response
 
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
 chatbot = gr.ChatInterface(
     respond,
     type="messages",
@@ -61,13 +81,8 @@ chatbot = gr.ChatInterface(
             label="Top-p (nucleus sampling)",
         ),
     ],
+    title="Llama 3.1 8B Tuned Chat"
 )
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
-
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()