fix?: history fmt

README.md

@@ -9,5 +9,5 @@ app_file: app.py
 pinned: false
 short_description: Reason about papers using LLMs
 license: agpl-3.0
-models: [HuggingFaceTB/SmolLM2-135M-Instruct]
+models: [HuggingFaceTB/SmolLM2-1.7B-Instruct]
 ---

app.py

@@ -1,3 +1,4 @@
+import os
 import time
 from typing import Dict, List, Optional, TypeAlias
 
@@ -7,11 +8,13 @@ import weave
 from papersai.utils import load_paper_as_context
 from transformers import pipeline
 
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 HistoryType: TypeAlias = List[Dict[str, str]]
 
 # Initialize the LLM and Weave client
 client = weave.init("papersai")
-checkpoint: str = "HuggingFaceTB/SmolLM2-135M-Instruct"
+checkpoint: str = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
 pipe = pipeline(
     model=checkpoint,
     torch_dtype=torch.bfloat16,
@@ -64,11 +67,14 @@ def invoke(history: HistoryType):
     Returns:
         BaseMessage: Response from the model
     """
-    input_text = pipe.tokenizer.apply_chat_template(history, tokenize=False)
-    response = pipe(input_text, do_sample=True, top_p=0.95, max_new_tokens=1024)[0][
+    input_text = pipe.tokenizer.apply_chat_template(
+        history + [{"role": "assistant", "content": f"Context: {state.context}\n"}],
+        tokenize=False,
+    )
+    response = pipe(input_text, do_sample=True, top_p=0.95, max_new_tokens=100)[0][
         "generated_text"
     ]
-    response = response.split("<|im_start|>assistant\n")[-1].split("<|im_end|>")[1]
+    response = response.split("\nassistant\n")[-1]
     return response