Update app.py

app.py

@@ -130,7 +130,10 @@ async def generate(request: GenerateRequest):
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay):
-    encoded_input = tokenizer(input_text, return_tensors="pt").to(device)
+    # Get the maximum model input length
+    max_model_length = model.config.max_position_embeddings
+
+    encoded_input = tokenizer(input_text, return_tensors="pt", max_length=max_model_length, truncation=True).to(device)
 
     def stop_criteria(input_ids, scores):
         decoded_output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
@@ -144,20 +147,25 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
     token_buffer = []
     output_ids = encoded_input.input_ids
     while True:
-        outputs = model.generate(
-            output_ids,
-            do_sample=generation_config.do_sample,
-            max_new_tokens=generation_config.max_new_tokens,
-            temperature=generation_config.temperature,
-            top_p=generation_config.top_p,
-            top_k=generation_config.top_k,
-            repetition_penalty=generation_config.repetition_penalty,
-            num_return_sequences=generation_config.num_return_sequences,
-            stopping_criteria=stopping_criteria,
-            output_scores=True,
-            return_dict_in_generate=True,
-            pad_token_id=tokenizer.pad_token_id
-        )
+        try:
+            outputs = model.generate(
+                output_ids,
+                do_sample=generation_config.do_sample,
+                max_new_tokens=generation_config.max_new_tokens,
+                temperature=generation_config.temperature,
+                top_p=generation_config.top_p,
+                top_k=generation_config.top_k,
+                repetition_penalty=generation_config.repetition_penalty,
+                num_return_sequences=generation_config.num_return_sequences,
+                stopping_criteria=stopping_criteria,
+                output_scores=True,
+                return_dict_in_generate=True,
+                pad_token_id=tokenizer.pad_token_id
+            )
+        except IndexError as e:
+            print(f"IndexError during generation: {e}")
+            break  # Exit the loop if there's an index error
+
         new_token_ids = outputs.sequences[0][encoded_input.input_ids.shape[-1]:]
 
         for token_id in new_token_ids: