fix?: system prompt

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-1.7B-Instruct]
+models: [HuggingFaceTB/SmolLM2-135M-Instruct]
 ---

app.py

@@ -14,7 +14,7 @@ HistoryType: TypeAlias = List[Dict[str, str]]
 
 # Initialize the LLM and Weave client
 client = weave.init("papersai")
-checkpoint: str = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
+checkpoint: str = "HuggingFaceTB/SmolLM2-135M-Instruct"
 pipe = pipeline(
     model=checkpoint,
     torch_dtype=torch.bfloat16,
@@ -68,7 +68,7 @@ def invoke(history: HistoryType):
         BaseMessage: Response from the model
     """
     input_text = pipe.tokenizer.apply_chat_template(
-        history + [{"role": "assistant", "content": f"Context: {state.context}\n"}],
+        history,
         tokenize=False,
     )
     response = pipe(input_text, do_sample=True, top_p=0.95, max_new_tokens=100)[0][
@@ -105,6 +105,9 @@ def update_state(history: HistoryType, message: Optional[Dict[str, str]]):
                 state.context = " ".join(doc_context)[
                     : pipe.model.config.max_position_embeddings
                 ]
+                history.append(
+                    {"role": "system", "content": f"Context: {state.context}\n"}
+                )
             except Exception as e:
                 history.append(
                     {"role": "assistant", "content": f"Error loading file: {str(e)}"}