fix?: slice based on max tokens

README.md

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

app.py

@@ -5,14 +5,13 @@ import gradio as gr
 import torch
 import weave
 from papersai.utils import load_paper_as_context
-from transformers import AutoTokenizer, pipeline
+from transformers import pipeline
 
 HistoryType: TypeAlias = List[Dict[str, str]]
 
 # Initialize the LLM and Weave client
 client = weave.init("papersai")
 checkpoint: str = "HuggingFaceTB/SmolLM2-135M-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 pipe = pipeline(
     model=checkpoint,
     torch_dtype=torch.bfloat16,
@@ -65,7 +64,7 @@ def invoke(history: HistoryType):
     Returns:
         BaseMessage: Response from the model
     """
-    input_text = tokenizer.apply_chat_template(history, tokenize=False)
+    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][
         "generated_text"
     ]
@@ -96,7 +95,14 @@ def update_state(history: HistoryType, message: Optional[Dict[str, str]]):
             try:
                 state.context = load_paper_as_context(file_path=file_path)
                 doc_context = [x.get_content() for x in state.context]
-                history.append({"role": "assistant", "content": " ".join(doc_context)})
+                history.append(
+                    {
+                        "role": "assistant",
+                        "content": [" ".join(doc_context)][
+                            : pipe.model.config.max_position_embeddings
+                        ],
+                    }
+                )
             except Exception as e:
                 history.append(
                     {"role": "assistant", "content": f"Error loading file: {str(e)}"}