Update app.py

app.py

@@ -3,15 +3,17 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 from threading import Thread
 
+# Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1")
 model = AutoModelForCausalLM.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1", torch_dtype=torch.float16)
-# Move model to GPU if available
+
+# Move model to GPU if available, otherwise use CPU
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = model.to(device)
 
 class StopOnTokens(StoppingCriteria):
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        stop_ids = [29, 0]
+        stop_ids = [29, 0]  # Define stop token IDs
         for stop_id in stop_ids:
             if input_ids[0][-1] == stop_id:
                 return True
@@ -21,11 +23,16 @@ def predict(message, history):
     history_transformer_format = list(zip(history[:-1], history[1:])) + [[message, ""]]
     stop = StopOnTokens()
 
-    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])
-                for item in history_transformer_format])
+    # Format the messages for the model
+    messages = "".join([f"\n<human>:{item[0]}\n<bot>:{item[1]}" for item in history_transformer_format])
 
-    model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
+    # Tokenize the input and move it to the correct device (GPU/CPU)
+    model_inputs = tokenizer([messages], return_tensors="pt").to(device)
+    
+    # Create a streamer for output token generation
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    
+    # Define generation parameters
     generate_kwargs = dict(
         model_inputs,
         streamer=streamer,
@@ -36,20 +43,25 @@ def predict(message, history):
         temperature=1.0,
         num_beams=1,
         stopping_criteria=StoppingCriteriaList([stop])
-        )
+    )
+
+    # Run the generation in a separate thread
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
+    # Yield generated tokens as they are produced
     partial_message = ""
     for new_token in streamer:
-        if new_token != '<':
+        if new_token != '<':  # Ignore special tokens
             partial_message += new_token
             yield partial_message
 
+# Gradio interface to interact with the model
 gr.ChatInterface(predict).launch()
 
 
 
+
 # import gradio as gr
 # from huggingface_hub import InferenceClient