Update app.py

app.py

@@ -27,19 +27,23 @@ def generate_roast(resume_text):
     # Define the prompt separately
     prompt_text = "Roast this resume:\n\n"
     
-    # Tokenize the prompt and resume text, truncating to fit within the 2048 token limit
-    max_tokens = 2048 - len(tokenizer(prompt_text)['input_ids'])  # Reserve space for the prompt
-    inputs = tokenizer(resume_text, return_tensors="pt", truncation=True, max_length=max_tokens)
-
-    # Convert the tokenized inputs back to text for the prompt
-    truncated_resume_text = tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True)
+    # Calculate how many tokens the prompt uses
+    prompt_tokens = tokenizer(prompt_text, return_tensors="pt")['input_ids'].shape[1]
+    
+    # Ensure the total length (prompt + resume text) doesn't exceed 2048 tokens
+    max_resume_tokens = 2048 - prompt_tokens
+    
+    # Tokenize the resume text and truncate to max_resume_tokens
+    resume_tokens = tokenizer(resume_text, return_tensors="pt", truncation=True, max_length=max_resume_tokens)
+    
+    # Decode the truncated resume back into text
+    truncated_resume_text = tokenizer.decode(resume_tokens['input_ids'][0], skip_special_tokens=True)
 
-    # Create the final prompt by combining the prompt and the truncated resume
+    # Create the final prompt with the truncated resume text
     prompt = f"{prompt_text}{truncated_resume_text}\n\nRoast:"
-
-    generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
     
     # Generate roast
+    generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
     roast = generator(prompt, max_new_tokens=50, num_return_sequences=1)
     
     return roast[0]['generated_text']
@@ -47,6 +51,7 @@ def generate_roast(resume_text):
 
 
 
+
 # Gradio interface function
 def roast_resume(file):
     if file.name.endswith('.pdf'):