Update app.py

app.py

@@ -5,24 +5,25 @@ import nltk
 import os
 import uvicorn
 
+from chunker import chunk_by_token_limit  
+
 nltk.download('punkt', quiet=True)
 
 app = FastAPI()
 
-HF_AUTH_TOKEN = os.getenv("HF_TOKEN") 
+HF_AUTH_TOKEN = os.getenv("HF_TOKEN")
 
 MODEL_NAME = "VincentMuriuki/legal-summarizer"
-summarizer = pipeline("summarization", model=MODEL_NAME, token=HF_AUTH_TOKEN)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_AUTH_TOKEN)
+summarizer = pipeline("summarization", model=MODEL_NAME, use_auth_token=HF_AUTH_TOKEN)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_AUTH_TOKEN)
 
 class SummarizeInput(BaseModel):
     text: str
 
 class ChunkInput(BaseModel):
     text: str
-    max_tokens: int = 512  # Default chunk size
+    max_tokens: int = 1024  
 
-# Summarize endpoint
 @app.post("/summarize")
 def summarize_text(data: SummarizeInput):
     summary = summarizer(data.text, max_length=150, min_length=30, do_sample=False)
@@ -30,25 +31,9 @@ def summarize_text(data: SummarizeInput):
 
 @app.post("/chunk")
 def chunk_text(data: ChunkInput):
-    sentences = nltk.sent_tokenize(data.text)
-    chunks = []
-    current_chunk = ""
-    current_token_count = 0
-
-    for sentence in sentences:
-        token_count = len(tokenizer.tokenize(sentence))
-        if current_token_count + token_count > data.max_tokens:
-            if current_chunk:
-                chunks.append(current_chunk.strip())
-            current_chunk = sentence
-            current_token_count = token_count
-        else:
-            current_chunk = f"{current_chunk} {sentence}".strip()
-            current_token_count += token_count
-
-    if current_chunk:
-        chunks.append(current_chunk.strip())
-
+    chunks = chunk_by_token_limit(data.text, data.max_tokens, tokenizer)
     return {"chunks": chunks}
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)
+