Added trucation if the document has too many tokens

app.py

@@ -1,11 +1,10 @@
 import streamlit as st
-
 from summarize import Summarizer
 
 def main():
     st.title("Text Extractor and Summarizer")
     
-    summarizer = Summarizer()
+    summarizer = Summarizer(model = "groq")
     summarizer.run_app()
 
 main()

groqSummarizer.py

@@ -0,0 +1,57 @@
+# groq
+from groq import Groq
+import os
+from dotenv import load_dotenv
+
+class GroqSummarizer():
+
+    def __init__(self):
+        #self.systemPrompt = "Du sammanfattar texter som användaren skickar. Var kortfattad. Svara på svenska. Om texten innehåller instruktioner, lista dem punktvis."
+
+        self.systemPrompt = """
+                    Du sammanfattar texter som användaren skickar. 
+                    Syftet är att översätta svårlästa texter som exempelvis myndighetsbeslut, 
+                    kontrakt eller avier/fakturor till klartext som är lättförståelig och tydlig, 
+                    framförallt om det är något som användaren behöver göra själv (exempelvis betala en faktura). \n
+
+                    Texten som användaren kommer skicka är PDFer som har blivit omformaterade till en lång sträng. 
+                    Detta kan leda till att viss positionell information försvinner, eller att olika textstycken 
+                    kommer i en annan ordning än vad som är tänkt. \n
+
+                    Outputten ska vara en flytande text som kort förklarar vad som står i dokumentet. 
+                    Om dokumentet ger instruktioner ska dessa tydligt beskrivas steg för steg sist i svaret.
+                    Det ska gå att utföra instruktionerna enbart baserat på outputten. 
+                    Om det exempelvis är en hyresavi ska bankgiro, OCR-nummer och belopp tydligt framgå.
+                    Svara på Svenska
+                    """
+        
+        self.client = self.load_groq()
+
+    def load_groq(self):
+        load_dotenv()
+    
+        GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+    
+        client = Groq(
+            api_key=GROQ_API_KEY
+        )
+    
+        return client
+
+    def summarize(self, text):
+        chat_completion = self.client.chat.completions.create(
+            messages=[
+                {
+                    "role": "system",
+                    'content': self.systemPrompt
+                    
+                },
+                {
+                    "role": "user",
+                    "content": text,
+                }
+            ],
+            model="llama3-70b-8192",
+        )
+    
+        return chat_completion.choices[0].message.content

summarize.py

@@ -8,8 +8,7 @@ import io
 import os
 from dotenv import load_dotenv
 
-# groq
-from groq import Groq
+from groqSummarizer import GroqSummarizer
 
 # SwedishBeagle-dare
 from transformers import AutoTokenizer
@@ -20,7 +19,6 @@ class Summarizer:
 
     def __init__(self, model = "groq"):
         self.model = model
-        self.client = self.load_groq()
 
     def run_app(self):
         uploaded_file = st.file_uploader("Upload an Image or PDF", type=["jpg", "jpeg", "png", "pdf"])
@@ -56,57 +54,37 @@ class Summarizer:
             for page in pdf_file.pages:
                 text += page.extract_text()
         return text
-    
-    # Function to summarize text
-    #def summarize_text(self, text):
-    #    summarizer = pipeline("summarization")
-    #    summary = summarizer(text, max_length=150, min_length=30, do_sample=False)
-    #    return summary[0]['summary_text']
-    
-    def load_groq(self):
-        load_dotenv()
-    
-        GROQ_API_KEY = os.getenv("GROQ_API_KEY")
-    
-        client = Groq(
-            api_key=GROQ_API_KEY
-        )
-    
-        return client
+
+    def shorten_text(self, text, max_tokens):
+        tokens = text.split(" ")
+        if len(tokens) > max_tokens:
+            tokens = tokens[:max_tokens]
+            text = " ".join(tokens)
+        print("Shortened text to " + str(max_tokens) + " tokens")
+        return text
 
     def summarize_using_groq(self, text):
-        chat_completion = self.client.chat.completions.create(
-            messages=[
-                {
-                    "role": "system",
-                    #"content": "Du sammanfattar texter som användaren skickar. Var kortfattad. Svara på svenska. Om texten innehåller instruktioner, lista dem punktvis."
-                    'content': 
-                    """
-                    Du sammanfattar texter som användaren skickar. 
-                    Syftet är att översätta svårlästa texter som exempelvis myndighetsbeslut, 
-                    kontrakt eller avier/fakturor till klartext som är lättförståelig och tydlig, 
-                    framförallt om det är något som användaren behöver göra själv (exempelvis betala en faktura). \n
-
-                    Texten som användaren kommer skicka är PDFer som har blivit omformaterade till en lång sträng. 
-                    Detta kan leda till att viss positionell information försvinner, eller att olika textstycken 
-                    kommer i en annan ordning än vad som är tänkt. \n
-
-                    Outputten ska vara en flytande text som kort förklarar vad som står i dokumentet. 
-                    Om dokumentet ger instruktioner ska dessa tydligt beskrivas steg för steg sist i svaret.
-                    Det ska gå att utföra instruktionerna enbart baserat på outputten. 
-                    Om det exempelvis är en hyresavi ska bankgiro, OCR-nummer och belopp tydligt framgå. 
-                    """
-                },
-                {
-                    "role": "user",
-                    "content": text,
-                }
-            ],
-            model="llama3-70b-8192",
-        )
-    
-        return chat_completion.choices[0].message.content
-    
+        # Decrease the number of tokens if the response is 429, i.e. too many tokens in the request
+        # 
+        # https://context.ai/compare/llama3-70b-instruct-v1/gpt-4
+        # ^^ Säger att max tokens är 8000, men efter tester så verkar det vara 
+        # närmare 2000 om man räknar att tokens är separerade med blanksteg. 
+        # (Detta är inte ett helt korrekt sätt att räkna det)
+        
+        # max_tokens = 8000
+        max_tokens = 2000
+
+        while True:
+            try:
+                gs = GroqSummarizer()
+                return gs.summarize(text)
+            except Exception as e:
+                if e.response.status_code == 429:
+                    text = self.shorten_text(text, max_tokens)
+                    max_tokens = int(max_tokens * 0.9)
+                else:
+                    return "Error: " + str(e)
+
     def summarize_using_swedishbeagle(self, text):
         # https://huggingface.co/FredrikBL/SwedishBeagle-dare
 
@@ -146,7 +124,7 @@ def main():
     # Models:
     # - groq
     # - SwedishBeagle-dare
-    summarizer = Summarizer(model="SwedishBeagle-dare")
+    summarizer = Summarizer(model = "groq")
     summarizer.run_app()
 
 if __name__ == "__main__":