Create app.py

app.py

@@ -0,0 +1,195 @@
+import os
+import gradio as gr
+import pdfplumber
+import requests
+import faiss
+import json
+import torch
+from bs4 import BeautifulSoup
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from sentence_transformers import SentenceTransformer
+import numpy as np
+import tempfile
+import logging
+from datetime import datetime
+from typing import List, Dict
+
+# Optimize CUDA memory management
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class CaseStudyGenerator:
+    def __init__(self):
+        self.model_name = "facebook/opt-2.7b"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Clear any reserved memory
+        if self.device == "cuda":
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+        model_kwargs = {
+            'torch_dtype': torch.float16 if self.device == "cuda" else torch.float32
+        }
+
+        try:
+            self.model = AutoModelForCausalLM.from_pretrained(self.model_name, **model_kwargs)
+            if self.device == "cuda":
+                self.model = self.model.to(self.device)
+                self.model.gradient_checkpointing_enable()
+        except RuntimeError as e:
+            logger.warning(f"Memory issue detected: {e}, attempting 8-bit loading.")
+
+            try:
+                from transformers import BitsAndBytesConfig
+                quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+                self.model = AutoModelForCausalLM.from_pretrained(self.model_name, quantization_config=quantization_config)
+            except ImportError:
+                logger.error("Missing 'bitsandbytes'. Install it using 'pip install -U bitsandbytes'")
+                logger.info("Switching to CPU to continue operations.")
+                self.device = "cpu"
+                self.model = AutoModelForCausalLM.from_pretrained(self.model_name, torch_dtype=torch.float32)
+
+        self.generator = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            device=0 if self.device == "cuda" else -1,
+            max_length=2048,
+            num_return_sequences=1,
+            temperature=0.8,
+            top_p=0.95,
+            do_sample=True,
+            pad_token_id=self.tokenizer.eos_token_id
+        )
+
+        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+        self.dimension = 384
+        self.index = faiss.IndexFlatL2(self.dimension)
+        self.stored_texts: List[Dict] = []
+
+    def clean_url(self, url: str) -> str:
+        if not url.startswith(('http://', 'https://')):
+            return ""
+        return url.split('?')[0][:100]
+
+    def fetch_articles(self, topic: str) -> List[str]:
+        try:
+            search_url = f"https://www.google.com/search?q={topic.replace(' ', '+')}+case+study+manufacturing+strategy"
+            headers = {"User-Agent": "Mozilla/5.0"}
+            response = requests.get(search_url, headers=headers, timeout=10)
+            response.raise_for_status()
+
+            soup = BeautifulSoup(response.text, "html.parser")
+            articles = [self.clean_url(link.get("href", "")) for link in soup.find_all("a") if "google" not in link.get("href", "")]
+            return articles[:5] or ["No articles found"]
+        except Exception as e:
+            logger.error(f"Error fetching articles: {str(e)}")
+            return ["Error fetching articles"]
+
+    def process_pdf(self, pdf_file) -> str:
+        try:
+            if pdf_file is None:
+                return "No PDF provided"
+
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_pdf:
+                temp_pdf.write(pdf_file.read())
+                temp_path = temp_pdf.name
+
+            text = []
+            with pdfplumber.open(temp_path) as pdf:
+                text = [page.extract_text().strip() for page in pdf.pages if page.extract_text()]
+
+            os.unlink(temp_path)
+            return "\n".join(text) or "No text extracted from PDF"
+        except Exception as e:
+            logger.error(f"Error processing PDF: {str(e)}")
+            return "Error processing PDF"
+
+    def generate_case_study(self, topic: str, pdf=None) -> str:
+        try:
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+
+            articles = self.fetch_articles(topic)
+            pdf_text = self.process_pdf(pdf) if pdf else "No PDF provided"
+
+            prompt = f"""Write a professional case study about {topic}.
+Background Information:
+- Topic: {topic}
+- Supporting Documents: {pdf_text[:500]}
+- Related Sources: {', '.join(articles)}
+
+Format your response as:
+1. Executive Summary
+2. Company Background
+3. Challenge Analysis
+4. Strategic Implementation
+5. Results and Impact
+6. Key Learnings
+"""
+
+            output = self.generator(
+                prompt,
+                max_new_tokens=1024,
+                num_return_sequences=1,
+                temperature=0.8,
+                top_p=0.95,
+                do_sample=True,
+                repetition_penalty=1.2,
+                no_repeat_ngram_size=3
+            )
+
+            case_study = output[0]['generated_text'].replace(prompt, "").strip()
+            embedding = self.embedding_model.encode([case_study])[0]
+            self.index.add(embedding.reshape(1, -1))
+
+            self.stored_texts.append({
+                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                "topic": topic,
+                "content": case_study
+            })
+
+            return case_study
+        except Exception as e:
+            logger.error(f"Error generating case study: {str(e)}")
+            return f"Error generating case study: {str(e)}"
+
+    def retrieve_past_case_studies(self) -> str:
+        try:
+            if not self.stored_texts:
+                return "No case studies generated yet."
+
+            result = ""
+            for idx, case in enumerate(self.stored_texts[-5:], start=1):
+                result += f"Case Study {idx}\nTopic: {case['topic']}\nGenerated on: {case['timestamp']}\n\n{case['content']}\n\n=== End of Case Study ===\n\n"
+            return result
+        except Exception as e:
+            logger.error(f"Error retrieving past case studies: {str(e)}")
+            return "Error retrieving past case studies"
+
+# Gradio interface
+with gr.Blocks() as app:
+    gr.Markdown("# AI Case Study Generator (Optimized for GPU-T4 & CPU)")
+    with gr.Row():
+        topic = gr.Textbox(label="Enter Topic")
+        pdf = gr.File(label="Upload PDF", type="binary")
+    with gr.Row():
+        generate_btn = gr.Button("Generate Case Study")
+        retrieve_btn = gr.Button("Retrieve Past Case Studies")
+    output = gr.Textbox(label="Generated Case Study", lines=20)
+    past_cases = gr.Textbox(label="Past Case Studies", lines=20)
+
+    generator = CaseStudyGenerator()
+    generate_btn.click(generator.generate_case_study, inputs=[topic, pdf], outputs=output)
+    retrieve_btn.click(generator.retrieve_past_case_studies, outputs=past_cases)
+
+# Launch the application
+if __name__ == "__main__":
+    app.launch(share=True)  # Remove enable_queue
+    # or, If using Gradio 3.x or later, use:
+    # app.queue().launch(share=True)