Update app.py

app.py

@@ -6,6 +6,7 @@ from threading import Thread
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
 model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
+model = model.to('cuda:0')
 
 class StopOnTokens(StoppingCriteria):
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
@@ -15,7 +16,7 @@ class StopOnTokens(StoppingCriteria):
                 return True
         return False
 
-def predict(message, history, temperature, max_new_tokens, min_new_tokens):
+def predict(message, history):
 
     history_transformer_format = history + [[message, ""]]
     stop = StopOnTokens()
@@ -23,15 +24,18 @@ def predict(message, history, temperature, max_new_tokens, min_new_tokens):
     messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])  #curr_system_message +
                 for item in history_transformer_format])
 
-    model_inputs = tokenizer([messages], return_tensors="pt")
+    model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
+    temperature = gr.Slider(minimum=0.1, maximum=2.0, value=1.0, label="Temperature"),
+    max_new_tokens = gr.Slider(minimum=0, maximum=2048, value=10, label="Temperature"),
+    min_new_tokens = gr.Slider(minimum=0, maximum=2048, value=1, label="Temperature"),
     generate_kwargs = dict(
         model_inputs,
-        max_new_tokens=int(max_new_tokens),
-        min_new_tokens=int(min_new_tokens),
+        max_new_tokens=int(max_new_tokens.value),
+        min_new_tokens=int(min_new_tokens.value),
         do_sample=True,
         top_p=1,
         top_k=50,
-        temperature=float(temperature),
+        temperature=float(temperature.value),
         num_beams=1,
         stopping_criteria=StoppingCriteriaList([stop])
         )
@@ -45,16 +49,4 @@ def predict(message, history, temperature, max_new_tokens, min_new_tokens):
             yield partial_message
 
 
-iface = gr.Interface(
-    fn=predict,
-    inputs=["text", "text", gr.Slider(minimum=0.1, maximum=2.0, default=1.0, label="Temperature"),
-            gr.Slider(minimum=1, maximum=2048, default=1024, label="Max Tokens"),
-            gr.Slider(minimum=1, maximum=1024, default=1, label="Min Tokens")],
-    outputs="text",
-    live=True,
-    capture_session=True,
-    layout="vertical",
-    chat=True
-)
-
-iface.launch()
+gr.ChatInterface(predict).queue().launch()