Update app.py

app.py

@@ -19,7 +19,8 @@ logger = logging.getLogger(__name__)
 
 # Constants
 EMBEDDING_MODEL = 'sentence-transformers/all-MiniLM-L6-v2'
-DEFAULT_MODEL = "distilgpt2"  # A smaller model that's more likely to work in Spaces
+DEFAULT_MODEL = "distilgpt2"
+DEFAULT_MAX_LENGTH = 1024  # Increased default max length
 
 # Check for GPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -36,12 +37,12 @@ def load_embeddings():
         return None
 
 @st.cache_resource
-def load_llm(model_name):
+def load_llm(model_name, max_length):
     """Load and cache the language model."""
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(model_name)
-        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device, max_length=512)
+        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device, max_length=max_length)
         return HuggingFacePipeline(pipeline=pipe)
     except Exception as e:
         logger.error(f"Failed to load LLM: {e}")
@@ -78,7 +79,7 @@ def summarize_report(documents: List[Document], llm) -> str:
     """Summarize the report using the loaded model."""
     try:
         prompt_template = """
-        Summarize the following text in a clear and concise manner:
+        Summarize the following text in a clear and concise manner. Focus on the main points and key details:
 
         {text}
 
@@ -99,10 +100,11 @@ def main():
     st.title("Report Summarizer")
 
     model_option = st.sidebar.text_input("Enter model name", value=DEFAULT_MODEL)
+    max_length = st.sidebar.slider("Max summary length", min_value=256, max_value=2048, value=DEFAULT_MAX_LENGTH, step=128)
 
     uploaded_file = st.sidebar.file_uploader("Upload your Report", type="pdf")
 
-    llm = load_llm(model_option)
+    llm = load_llm(model_option, max_length)
     if not llm:
         st.error(f"Failed to load the model {model_option}. Please try another model.")
         return