Update app.py

app.py

@@ -27,35 +27,30 @@ def generate_roast(resume_text):
     # Define the prompt separately
     prompt_text = "Roast this resume:\n\n"
     
-    # Tokenize the prompt to calculate its token length
+    # Tokenize the prompt
     prompt_tokenized = tokenizer(prompt_text, return_tensors="pt")
     prompt_tokens = prompt_tokenized['input_ids'].shape[1]
     
-    # Calculate the remaining tokens for the resume text (2048 - prompt tokens)
+    # Calculate remaining tokens for resume text (2048 - prompt tokens)
     max_resume_tokens = 2048 - prompt_tokens
     
-    # Tokenize and truncate the resume text to fit within the remaining token limit
-    resume_tokenized = tokenizer(resume_text, truncation=True, max_length=max_resume_tokens)
+    # Tokenize and strictly truncate the resume text to fit within the remaining token space
+    resume_tokenized = tokenizer(resume_text, truncation=True, max_length=max_resume_tokens, return_tensors="pt")
     
     # Decode the truncated resume back into a string
-    truncated_resume_text = tokenizer.decode(resume_tokenized['input_ids'], skip_special_tokens=True)
+    truncated_resume_text = tokenizer.decode(resume_tokenized['input_ids'][0], skip_special_tokens=True)
     
     # Combine the prompt and the truncated resume text
     final_prompt = f"{prompt_text}{truncated_resume_text}\n\nRoast:"
     
-    # Generate roast with the truncated prompt
+    # Generate the roast
     generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
     
-    # Generate roast within max_new_tokens limit
+    # Ensure generated roast doesn't exceed token limit
     roast = generator(final_prompt, max_new_tokens=50, num_return_sequences=1)
     
     return roast[0]['generated_text']
 
-
-
-
-
-
 # Gradio interface function
 def roast_resume(file):
     if file.name.endswith('.pdf'):