a

app.py

@@ -5,8 +5,9 @@ from huggingface_hub import InferenceClient
 Copied from inference in colab notebook
 """
 
-from transformers import AutoModel, AutoTokenizer , AutoModelForCausalLM
+from transformers import AutoTokenizer , AutoModelForCausalLM , TextIteratorStreamer
 import torch
+from threading import Thread
 
 # Load model and tokenizer globally to avoid reloading for every request
 model_path = "Heit39/llama_lora_model_1"
@@ -23,6 +24,58 @@ model = PeftModel.from_pretrained(base_model, model_path)
 
 
 # Define the response function
+# def respond(
+#     message: str,
+#     history: list[tuple[str, str]],
+#     system_message: str,
+#     max_tokens: int,
+#     temperature: float,
+#     top_p: float,
+# ):
+#     # Combine system message and history into a single prompt
+#     messages = [{"role": "system", "content": system_message}]
+#     for val in history:
+#         if val[0]:
+#             messages.append({"role": "user", "content": val[0]})
+#         if val[1]:
+#             messages.append({"role": "assistant", "content": val[1]})
+#     messages.append({"role": "user", "content": message})
+    
+#     # Create a single text prompt from the messages
+#     prompt = ""
+#     for msg in messages:
+#         if msg["role"] == "system":
+#             prompt += f"[System]: {msg['content']}\n\n"
+#         elif msg["role"] == "user":
+#             prompt += f"[User]: {msg['content']}\n\n"
+#         elif msg["role"] == "assistant":
+#             prompt += f"[Assistant]: {msg['content']}\n\n"
+    
+#     # Tokenize the prompt
+#     inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+#     input_ids = inputs.input_ids.to("cpu")  # Ensure input is on the CPU
+
+#     # Generate response
+#     output_ids = model.generate(
+#         input_ids,
+#         max_length=input_ids.shape[1] + max_tokens,
+#         temperature=temperature,
+#         top_p=top_p,
+#         do_sample=True,
+#     )
+    
+#     # Decode the generated text
+#     generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+#     # Extract the assistant's response from the generated text
+#     assistant_response = generated_text[len(prompt):].strip()
+
+#     # Yield responses incrementally (simulate streaming)
+#     response = ""
+#     for token in assistant_response.split():  # Split tokens by whitespace
+#         response += token + " "
+#         yield response.strip()
+
 def respond(
     message: str,
     history: list[tuple[str, str]],
@@ -54,27 +107,26 @@ def respond(
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
     input_ids = inputs.input_ids.to("cpu")  # Ensure input is on the CPU
 
-    # Generate response
-    output_ids = model.generate(
-        input_ids,
-        max_length=input_ids.shape[1] + max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=True,
-    )
-    
-    # Decode the generated text
-    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    # Generate tokens incrementally
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    generation_kwargs = {
+        "input_ids": input_ids,
+        "max_new_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "do_sample": True,
+        "streamer": streamer,
+    }
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
 
-    # Extract the assistant's response from the generated text
-    assistant_response = generated_text[len(prompt):].strip()
-
-    # Yield responses incrementally (simulate streaming)
+    # Yield responses as they are generated
     response = ""
-    for token in assistant_response.split():  # Split tokens by whitespace
-        response += token + " "
+    for token in streamer:
+        response += token
         yield response.strip()
 
+
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """