Update app.py

app.py

@@ -1,47 +1,33 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
-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")
 
-class StopOnTokens(StoppingCriteria):
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        stop_ids = [29, 0]
-        for stop_id in stop_ids:
-            if input_ids[0][-1] == stop_id:
-                return True
-        return False
+def generate_text(prompt, temperature, max_length, min_length):
+    # Tokenize the prompt
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
 
-def predict(message, history):
+    # Generate text using the model
+    output = model.generate(input_ids, max_length=max_length, min_length=min_length, temperature=temperature, num_return_sequences=1)
 
-    history_transformer_format = history + [[message, ""]]
-    stop = StopOnTokens()
+    # Decode the generated output
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    
+    return generated_text
 
-    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])  #curr_system_message +
-                for item in history_transformer_format])
+def chatbot_app(prompt, temperature, max_length, min_length):
+    generated_text = generate_text(prompt, temperature, max_length, min_length)
+    return generated_text
 
-    model_inputs = tokenizer([messages], return_tensors="pt")
-    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.value),
-        min_new_tokens=int(min_new_tokens.value),
-        do_sample=True,
-        top_p=1,
-        top_k=50,
-        temperature=float(temperature.value),
-        num_beams=1,
-        stopping_criteria=StoppingCriteriaList([stop])
-        )
-    generated_sequence = model.generate(**generate_kwargs)[0]
-    generated_text = tokenizer.decode(generated_sequence, skip_special_tokens=True)
+iface = gr.Interface(
+    fn=chatbot_app,
+    inputs=["text", gr.Number(minimum=0.1, maximum=2.0, value=1.0, label="Temperature"),
+            gr.Number(minimum=10, maximum=2048, value=10, label="Max Length"),
+            gr.Number(minimum=1, maximum=2048, value=1, label="Min Length")],
+    outputs="text",
+    live=False,
+)
 
-    yield generated_text
-
-
-gr.ChatInterface(predict).queue().launch()
+iface.launch()