Update app.py

app.py

@@ -7,26 +7,30 @@ from threading import Thread
 torch.set_num_threads(2)
 
 # Loading the tokenizer and model from Hugging Face's model hub.
-tokenizer = AutoTokenizer.from_pretrained("cnmoro/jack-68m-text-structurization")
-model = AutoModelForCausalLM.from_pretrained("cnmoro/jack-68m-text-structurization")
+tokenizer = AutoTokenizer.from_pretrained("cnmoro/teenytinyllama-160m-text-simplification-ptbr")
+model = AutoModelForCausalLM.from_pretrained("cnmoro/teenytinyllama-160m-text-simplification-ptbr")
 
 # using CUDA for an optimal experience
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = model.to(device)
 
+def count_tokens(text):
+    return len(tokenizer.tokenize(text))
+
 # Function to generate model predictions.
 def predict(message, history):
-    
+
+    formatted_prompt = f"<s><system>O objetivo é comprimir e estruturar o texto a seguir<texto>{message}</texto>"
     model_inputs = tokenizer([
-        f"### Structurize: {message}\n\n### Response:\n"
+        formatted_prompt
     ], return_tensors="pt").to(device)
     
-    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(tokenizer, timeout=120., skip_prompt=True, skip_special_tokens=True)
     
     generate_kwargs = dict(
         model_inputs,
         streamer=streamer,
-        max_new_tokens=512,
+        max_new_tokens=3072 - count_tokens(formatted_prompt),
         top_p=0.2,
         top_k=20,
         temperature=0.1,
@@ -43,6 +47,6 @@ def predict(message, history):
 
 # Setting up the Gradio chat interface.
 gr.ChatInterface(predict,
-                 title="TextStructurization_Jack68m_CPU",
+                 title="TextStructurization_TeenyTinyLlama160m_CPU",
                  description="Pass a text to be structurized"
                  ).launch()  # Launching the web interface.