Update app.py

app.py

@@ -78,163 +78,165 @@ def clear_gpu_memory():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
 
-# Let the user specify the column name for tweets text (defaulting to "content")
-tweets_column = st.text_input("Enter the column name for Tweets:", value="content")
-
-# Input: Query question for scoring and CSV file upload for candidate tweets
-query_input = st.text_area("Enter your query question for analysis (this does not need to be part of the CSV):")
-uploaded_file = st.file_uploader(f"Upload Tweets CSV File (must contain a '{tweets_column}' column)", type=["csv"])
-
-candidate_docs = []
-if uploaded_file is not None:
-    try:
-        df = pd.read_csv(uploaded_file)
-        if tweets_column not in df.columns:
+# Main Function Part:
+def main():
+    # Let the user specify the column name for tweets text (defaulting to "content")
+    tweets_column = st.text_input("Enter the column name for Tweets:", value="content")
+    
+    # Input: Query question for scoring and CSV file upload for candidate tweets
+    query_input = st.text_area("Enter your query question for analysis (this does not need to be part of the CSV):")
+    uploaded_file = st.file_uploader(f"Upload Tweets CSV File (must contain a '{tweets_column}' column)", type=["csv"])
+    
+    candidate_docs = []
+    if uploaded_file is not None:
+        try:
+            df = pd.read_csv(uploaded_file)
+            if tweets_column not in df.columns:
+                st.error(f"CSV must contain a '{tweets_column}' column.")
+            else:
+                candidate_docs = df[tweets_column].dropna().astype(str).tolist()
+        except Exception as e:
+            st.error(f"Error reading CSV file: {e}")
+    
+    if st.button("Generate Report"):
+        # Reset timer state so that the timer always shows up
+        st.session_state.timer_started = False
+        st.session_state.timer_frozen = False
+        if uploaded_file is None:
+            st.error("Please upload a CSV file.")
+        elif not tweets_column.strip():
+            st.error("Please enter your column name")
+        elif not candidate_docs:
             st.error(f"CSV must contain a '{tweets_column}' column.")
+        elif not query_input.strip():
+            st.error("Please enter a query question!")
         else:
-            candidate_docs = df[tweets_column].dropna().astype(str).tolist()
-    except Exception as e:
-        st.error(f"Error reading CSV file: {e}")
-
-if st.button("Generate Report"):
-    # Reset timer state so that the timer always shows up
-    st.session_state.timer_started = False
-    st.session_state.timer_frozen = False
-    if uploaded_file is None:
-        st.error("Please upload a CSV file.")
-    elif not tweets_column.strip():
-        st.error("Please enter your column name")
-    elif not candidate_docs:
-        st.error(f"CSV must contain a '{tweets_column}' column.")
-    elif not query_input.strip():
-        st.error("Please enter a query question!")
-    else:
-        if not st.session_state.timer_started and not st.session_state.timer_frozen:
-            st.session_state.timer_started = True
-            html(timer(), height=50)
-        status_text = st.empty()
-        progress_bar = st.progress(0)
-        
-        
-        processed_docs = []
-        scored_results = []
-        
-        # First, check which documents need summarization
-        docs_to_summarize = []
-        docs_indices = []
-        
-        for i, doc in enumerate(candidate_docs):
-            if len(doc) > 280:
-                docs_to_summarize.append(doc)
-                docs_indices.append(i)
-        
-        # If we have documents to summarize, load Llama model first
-        if docs_to_summarize:
-            status_text.markdown("**📝 Loading summarization model...**")
-            t5_pipe = get_summary_model()
+            if not st.session_state.timer_started and not st.session_state.timer_frozen:
+                st.session_state.timer_started = True
+                html(timer(), height=50)
+            status_text = st.empty()
+            progress_bar = st.progress(0)
             
-            status_text.markdown("**📝 Summarizing long tweets...**")
             
-            # Process documents that need summarization
-            for idx, (i, doc) in enumerate(zip(docs_indices, docs_to_summarize)):
-                progress = int((idx / len(docs_to_summarize)) * 25)  # First quarter of progress
-                progress_bar.progress(progress)
+            processed_docs = []
+            scored_results = []
+            
+            # First, check which documents need summarization
+            docs_to_summarize = []
+            docs_indices = []
+            
+            for i, doc in enumerate(candidate_docs):
+                if len(doc) > 280:
+                    docs_to_summarize.append(doc)
+                    docs_indices.append(i)
+            
+            # If we have documents to summarize, load Llama model first
+            if docs_to_summarize:
+                status_text.markdown("**📝 Loading summarization model...**")
+                t5_pipe = get_summary_model()
+                
+                status_text.markdown("**📝 Summarizing long tweets...**")
+                
+                # Process documents that need summarization
+                for idx, (i, doc) in enumerate(zip(docs_indices, docs_to_summarize)):
+                    progress = int((idx / len(docs_to_summarize)) * 25)  # First quarter of progress
+                    progress_bar.progress(progress)
+                    
+                    input_text = "summarize: " + doc
+                    
+                    try:
+                        summary_result = t5_pipe(
+                            input_text, 
+                            max_length=128,
+                            min_length=10,
+                            no_repeat_ngram_size=2,
+                            num_beams=4,
+                            early_stopping=True,
+                            truncation=True
+                        )
+                        
+                        # Store the summary in place of the original text
+                        candidate_docs[i] = summary_result[0]['generated_text']
+                        
+                    except Exception as e:
+                        st.warning(f"Error summarizing document {i}: {str(e)}")
                 
-                input_text = "summarize: " + doc
+                # Clear Llama model from memory
+                del t5_pipe
+                import gc
+                gc.collect()
+                torch.cuda.empty_cache()
+            
+            # Now load sentiment model
+            status_text.markdown("**🔍 Loading sentiment analysis model...**")
+            progress_bar.progress(25)
+            score_pipe = get_sentiment_model()
+            
+            status_text.markdown("**🔍 Scoring documents...**")
+            
+            # Process each document with sentiment analysis
+            for i, doc in enumerate(candidate_docs):
+                progress_offset = 25 if docs_to_summarize else 0
+                progress = progress_offset + int((i / len(candidate_docs)) * (50 - progress_offset))
+                progress_bar.progress(progress)
                 
                 try:
-                    summary_result = t5_pipe(
-                        input_text, 
-                        max_length=128,
-                        min_length=10,
-                        no_repeat_ngram_size=2,
-                        num_beams=4,
-                        early_stopping=True,
-                        truncation=True
-                    )
+                    # Process with sentiment analysis
+                    result = score_pipe(doc, truncation=True, max_length=512)
                     
-                    # Store the summary in place of the original text
-                    candidate_docs[i] = summary_result[0]['generated_text']
+                    # If it's a list, get the first element
+                    if isinstance(result, list):
+                        result = result[0]
+                        
+                    processed_docs.append(doc)
+                    scored_results.append(result)
                     
                 except Exception as e:
-                    st.warning(f"Error summarizing document {i}: {str(e)}")
+                    st.warning(f"Error scoring document {i}: {str(e)}")
+                    processed_docs.append("Error processing this document")
+                    scored_results.append({"label": "NEUTRAL", "score": 1})
+                
+                # Display occasional status updates
+                if i % max(1, len(candidate_docs) // 10) == 0:
+                    status_text.markdown(f"**🔍 Scoring documents... ({i}/{len(candidate_docs)})**")
             
-            # Clear Llama model from memory
-            del t5_pipe
+            # Pair documents with scores
+            scored_docs = list(zip(processed_docs, [result.get("score", 0.5) for result in scored_results]))
+            
+            # Clear sentiment model from memory
+            del score_pipe
             import gc
             gc.collect()
             torch.cuda.empty_cache()
-        
-        # Now load sentiment model
-        status_text.markdown("**🔍 Loading sentiment analysis model...**")
-        progress_bar.progress(25)
-        score_pipe = get_sentiment_model()
-        
-        status_text.markdown("**🔍 Scoring documents...**")
-        
-        # Process each document with sentiment analysis
-        for i, doc in enumerate(candidate_docs):
-            progress_offset = 25 if docs_to_summarize else 0
-            progress = progress_offset + int((i / len(candidate_docs)) * (50 - progress_offset))
-            progress_bar.progress(progress)
             
-            try:
-                # Process with sentiment analysis
-                result = score_pipe(doc, truncation=True, max_length=512)
-                
-                # If it's a list, get the first element
-                if isinstance(result, list):
-                    result = result[0]
-                    
-                processed_docs.append(doc)
-                scored_results.append(result)
-                
-            except Exception as e:
-                st.warning(f"Error scoring document {i}: {str(e)}")
-                processed_docs.append("Error processing this document")
-                scored_results.append({"label": "NEUTRAL", "score": 1})
+            #print_gpu_status("After sentiment model deletion, VRAM")
             
-            # Display occasional status updates
-            if i % max(1, len(candidate_docs) // 10) == 0:
-                status_text.markdown(f"**🔍 Scoring documents... ({i}/{len(candidate_docs)})**")
-        
-        # Pair documents with scores
-        scored_docs = list(zip(processed_docs, [result.get("score", 0.5) for result in scored_results]))
-        
-        # Clear sentiment model from memory
-        del score_pipe
-        import gc
-        gc.collect()
-        torch.cuda.empty_cache()
-        
-        #print_gpu_status("After sentiment model deletion, VRAM")
-        
-        # Load Gemma for final report generation
-        status_text.markdown("**📊 Loading report generation model...**")
-        progress_bar.progress(67)
-        
-        # Make sure GPU memory is clear
-        clear_gpu_memory()
-        print_gpu_status("Before loading Gemma model, VRAM")
-        
-        # Set memory optimization environment variable
-        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-        
-        # Sample or summarize the data for Gemma to avoid memory issues
-        status_text.markdown("**📝 Preparing data for report generation...**")
-        progress_bar.progress(75)
-        
-        import random
-        max_tweets = 100
-        if len(scored_docs) > max_tweets:
-            sampled_docs = random.sample(scored_docs, max_tweets)
-            st.info(f"Sampling {max_tweets} out of {len(scored_docs)} tweets for report generation")
-        else:
-            sampled_docs = scored_docs
-        
-        # Build prompt
-        messages = [
-            {"role": "user", "content": f"""
+            # Load Gemma for final report generation
+            status_text.markdown("**📊 Loading report generation model...**")
+            progress_bar.progress(67)
+            
+            # Make sure GPU memory is clear
+            clear_gpu_memory()
+            print_gpu_status("Before loading Gemma model, VRAM")
+            
+            # Set memory optimization environment variable
+            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+            
+            # Sample or summarize the data for Gemma to avoid memory issues
+            status_text.markdown("**📝 Preparing data for report generation...**")
+            progress_bar.progress(75)
+            
+            import random
+            max_tweets = 1000
+            if len(scored_docs) > max_tweets:
+                sampled_docs = random.sample(scored_docs, max_tweets)
+                st.info(f"Sampling {max_tweets} out of {len(scored_docs)} tweets for report generation")
+            else:
+                sampled_docs = scored_docs
+            
+            # Build prompt
+            messages = [
+                {"role": "user", "content": f"""
 Generate a well-structured business report based on tweets from twitter/X with sentiment score (0: negative, 1: neutral, 2: positive) that answers Query Question and meets following Requirements. 
 **Requirements:**
 - Include an introduction, key insights, and a conclusion.
@@ -246,60 +248,64 @@ Generate a well-structured business report based on tweets from twitter/X with s
 **Tweets with sentiment score:**
 {sampled_docs}
 Please ensure the report is complete and reaches approximately 1000 words.
-        """}
-        ]
-
-        # Create a process function to avoid the Triton registration issue
-        def process_with_gemma(messages):
-            # We'll define the pipeline here rather than using the cached version
-            # This ensures a clean library registration context
-            from transformers import pipeline, AutoTokenizer
-            import torch
-            
-            # Set dtype explicitly 
-            torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-            
-            try:
-                tokenizer = AutoTokenizer.from_pretrained("unsloth/gemma-3-1b-it")
-                pipe = pipeline(
-                    "text-generation",
-                    model="unsloth/gemma-3-1b-it",
-                    tokenizer=tokenizer,
-                    device=0 if torch.cuda.is_available() else -1,
-                    torch_dtype=torch_dtype
-                    
-                )
+                """}
+            ]
+    
+            # Create a process function to avoid the Triton registration issue
+            def process_with_gemma(messages):
+                # We'll define the pipeline here rather than using the cached version
+                # This ensures a clean library registration context
+                from transformers import pipeline, AutoTokenizer
+                import torch
                 
-                result = pipe(messages, max_new_tokens=1500, repetition_penalty=1.2, do_sample=True, temperature=0.7, return_full_text=False)
-                return result, None
+                # Set dtype explicitly 
+                torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
                 
-            except Exception as e:
-                return None, str(e)
-        
-        # Try to process with Gemma
-        status_text.markdown("**📝 Generating report with Gemma...**")
-        progress_bar.progress(80)
-        
-        raw_result, error = process_with_gemma(messages)
-
-        if error:
-                st.error(f"Gemma processing failed: {str(error)}")
-                report = "Error generating report. Please try again with fewer tweets."
-        else:
-            # Extract content from successful Gemma result
-            report = raw_result[0]['generated_text']
-            #extract_assistant_content(raw_result)
-        
-        progress_bar.progress(100)
-        status_text.success("**✅ Generation complete!**")
-        html("<script>localStorage.setItem('freezeTimer', 'true');</script>", height=0)
-        st.session_state.timer_frozen = True
+                try:
+                    tokenizer = AutoTokenizer.from_pretrained("unsloth/gemma-3-1b-it")
+                    pipe = pipeline(
+                        "text-generation",
+                        model="unsloth/gemma-3-1b-it",
+                        tokenizer=tokenizer,
+                        device=0 if torch.cuda.is_available() else -1,
+                        torch_dtype=torch_dtype
+                        
+                    )
+                    
+                    result = pipe(messages, max_new_tokens=1500, repetition_penalty=1.2, do_sample=True, temperature=0.7, return_full_text=False)
+                    return result, None
+                    
+                except Exception as e:
+                    return None, str(e)
+            
+            # Try to process with Gemma
+            status_text.markdown("**📝 Generating report with Gemma...**")
+            progress_bar.progress(80)
+            
+            raw_result, error = process_with_gemma(messages)
+    
+            if error:
+                    st.error(f"Gemma processing failed: {str(error)}")
+                    report = "Error generating report. Please try again with fewer tweets."
+            else:
+                # Extract content from successful Gemma result
+                report = raw_result[0]['generated_text']
+                #extract_assistant_content(raw_result)
+            
+            progress_bar.progress(100)
+            status_text.success("**✅ Generation complete!**")
+            html("<script>localStorage.setItem('freezeTimer', 'true');</script>", height=0)
+            st.session_state.timer_frozen = True
+    
+            # First, create the replacement separately
+            formatted_report = report.replace('\\n', '<br>')
+            
+            # Display title separately with standard formatting
+            st.subheader("Generated Report:")
+            
+            # Display the report content with normal styling
+            st.markdown(f"<div style='font-size: normal; font-weight: normal;'>{formatted_report}</div>", unsafe_allow_html=True)
 
-        # First, create the replacement separately
-        formatted_report = report.replace('\\n', '<br>')
-        
-        # Display title separately with standard formatting
-        st.subheader("Generated Report:")
-        
-        # Display the report content with normal styling
-        st.markdown(f"<div style='font-size: normal; font-weight: normal;'>{formatted_report}</div>", unsafe_allow_html=True)
+# Run the Main Function
+if __name__ == '__main__':
+    main()