Update app.py

app.py

@@ -5,27 +5,34 @@ import torch
 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
 model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
 
-def generate_text(prompt, temperature, max_length, min_length):
-    # Tokenize the prompt
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+def generate_text(system_prompt, user_prompt, temperature, max_length, min_length):
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt},
+    ]
+    tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
 
-    # Generate text using the model
-    output = model.generate(input_ids, max_length=max_length, min_length=min_length, temperature=temperature, num_return_sequences=1)
+    inputs = tokenizer(tokenized_chat, return_tensors="pt").to(model.device) 
+
+    outputs = model.generate(**inputs, use_cache=True, max_length=max_length, min_length=min_length, temperature=temperature, num_return_sequences=1)
+
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Decode the generated output
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    
     return generated_text
 
-def chatbot_app(prompt, temperature, max_length, min_length):
-    generated_text = generate_text(prompt, temperature, max_length, min_length)
+def chatbot_app(system_prompt, user_prompt, temperature, max_length, min_length):
+    generated_text = generate_text(system_prompt, user_prompt, temperature, max_length, min_length)
     return generated_text
 
 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")],
+    inputs=[
+        "text",
+        "text",  # Added a text box for the system prompt
+        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,
 )