content length

app.py

@@ -5,10 +5,13 @@ import os
 import os.path 
 import pickle
 import torch
+from openai import OpenAI
 
 base_model_id = "meta-llama/Llama-3.2-3B-Instruct"
 model_id = "HiGenius/Headline-Generation-Model"
+
 hf_token = os.environ.get('HF_TOKEN')
+openai_api_key = os.environ.get('OPENAI_API_KEY')
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -19,9 +22,23 @@ def load_model():
     tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_auth_token=hf_token)
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side='left'
+    tokenizer.truncation_side="left"
 
     return tokenizer, model
 
+def summarize_content(content):
+    client = OpenAI(api_key=openai_api_key)
+    response = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {"role": "system", "content": "Summarize the following article content concisely while preserving key information:"},
+            {"role": "user", "content": content}
+        ],
+        max_tokens=600,
+        temperature=0.3
+    )
+    return response.choices[0].message.content
+
 tokenizer, model = load_model()
 
 guideline_path = "./guidelines.txt"
@@ -29,6 +46,15 @@ with open(guideline_path, 'r', encoding='utf-8') as f:
     guidelines = f.read()
 
 def process_prompt(tokenizer, content, video_summary = '', guidelines = None):
+    # Check token lengths
+    content_tokens = len(tokenizer.encode(content))
+    total_tokens = content_tokens
+    if video_summary:
+        total_tokens += len(tokenizer.encode(video_summary))
+
+    if content_tokens > 850 or total_tokens > 900:
+        content = summarize_content(content)
+    
     if guidelines:
         system_prompt = "You are a helpful assistant that writes engaging headlines. To maximize engagement, you may follow these proven guidelines:\n" + guidelines
     else: