update sentence parsing, use model twice

app.py

@@ -17,15 +17,31 @@ model = PeftModel.from_pretrained(model, peft_model_id)
 
 
 def make_inference(original_text):
-    batch = tokenizer(
-        f"### Negative sentence:\n{original_text}\n\n### Reframing strategy: ['optimism']\n\n### Reframing sentence:\n",
+    str_strategy_prompt = f"### Negative sentence:\n{original_text}\n\n### Reframing strategy:\n"
+    batch_1 = tokenizer(
+        str_strategy_prompt,
         return_tensors="pt",
     )
 
     with torch.cuda.amp.autocast():
-        output_tokens = model.generate(**batch, max_new_tokens=50)
+        output_tokens_1 = model.generate(**batch_1, max_new_tokens=50)
 
-    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    output_1 = tokenizer.decode(output_tokens_1[0], skip_special_tokens=True)
+    reframing_strategy = output_1[len(str_strategy_prompt):].partition('\n')[0]
+
+    str_reframing_prompt = f"### Negative sentence:\n{original_text}\n\n### Reframing strategy:\n{reframing_strategy}\n\n### Reframing sentence:\n"
+    batch_2 = tokenizer(
+        str_reframing_prompt,
+        return_tensors="pt",
+    )
+
+    with torch.cuda.amp.autocast():
+        output_tokens_2 = model.generate(**batch_2, max_new_tokens=100)
+
+    output_2 = tokenizer.decode(output_tokens_2[0], skip_special_tokens=True)
+    reframing_sentence = output_2[len(str_reframing_prompt):]
+
+    return reframing_sentence
 
 
 if __name__ == "__main__":