Update app.py

app.py

@@ -69,14 +69,6 @@ def history_to_messages(history: List, system: str) -> List[Dict]:
         messages.append({'role': Role.ASSISTANT, 'content': h[1]})
     return messages
 
-def messages_to_history(messages: List[Dict]) -> Tuple[str, List]:
-    assert messages[0]['role'] == Role.SYSTEM
-    system = messages[0]['content']
-    history = []
-    for q, r in zip(messages[1::2], messages[2::2]):
-        history.append([q['content'], r['content']])
-    return system, history
-
 @spaces.GPU(duration=120)
 def generate(
     query: Optional[str],
@@ -97,26 +89,18 @@ def generate(
     messages = history_to_messages(history, system)
     messages.append({'role': Role.USER, 'content': query})
 
-    # Apply chat template and get input_ids
-    input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
-
-    # Create attention mask
-    attention_mask = torch.ones_like(input_ids)
-
-    # Trim input if it exceeds the maximum token length
-    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        attention_mask = attention_mask[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-
-    input_ids = input_ids.to(model.device)
-    attention_mask = attention_mask.to(model.device)
+    # Apply chat template and tokenize
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
     # Set up the streamer for real-time text generation
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
+        **model_inputs,
         streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=True,