Create app.py

app.py

@@ -0,0 +1,273 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from langchain_community.llms import HuggingFacePipeline
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+import warnings
+import os
+
+# Suppress warnings
+warnings.filterwarnings("ignore")
+os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
+
+# Model Configuration
+MODEL_NAME = "gpt2"
+EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+
+def initialize_models():
+    """Initialize language model and embedding model."""
+    try:
+        # Determine device
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Using device: {device}")
+
+        # Load model and tokenizer
+        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+        # Create pipeline
+        text_generation_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=512,
+            temperature=0.7,
+            repetition_penalty=1.1
+        )
+
+        # Langchain LLM
+        llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
+
+        # Embedding model
+        embedding_model = HuggingFaceEmbeddings(
+            model_name=EMBEDDING_MODEL,
+            model_kwargs={'device': str(device)}
+        )
+
+        return llm, embedding_model, model, tokenizer
+
+    except Exception as e:
+        print(f"Model initialization error: {e}")
+        return None, None, None, None
+
+# Initialize models
+llm, embedding_model, model, tokenizer = initialize_models()
+
+# Global variables for RAG state
+rag_retriever = None
+document_loaded = False
+loaded_doc_name = "No document loaded"
+
+def setup_rag_pipeline(doc_text, chunk_size=1000, chunk_overlap=150):
+    """Loads text, chunks, embeds, creates FAISS index, and sets up retriever."""
+    global rag_retriever, document_loaded, loaded_doc_name
+
+    if not doc_text or not isinstance(doc_text, str) or len(doc_text.strip()) == 0:
+        return "Error: No text provided or invalid input."
+
+    try:
+        # Text splitting
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            length_function=len,
+        )
+        docs = text_splitter.split_text(doc_text)
+
+        if not docs:
+            return "Error: Text splitting resulted in no documents."
+
+        # Create embeddings and FAISS index
+        vector_store = FAISS.from_texts(docs, embedding_model)
+        rag_retriever = vector_store.as_retriever(search_kwargs={"k": 3})
+
+        document_loaded = True
+        loaded_doc_name = f"Document processed ({len(doc_text)} chars, {len(docs)} chunks)."
+        return loaded_doc_name
+
+    except Exception as e:
+        document_loaded = False
+        rag_retriever = None
+        return f"Error processing document: {e}"
+
+def answer_question(question):
+    """Answers a question using the loaded RAG pipeline."""
+    if llm is None or embedding_model is None:
+        return "Error: Models not initialized properly."
+
+    if not document_loaded or rag_retriever is None:
+        return "Error: Please load a document before asking questions."
+
+    if not question or not isinstance(question, str) or len(question.strip()) == 0:
+        return "Error: Please enter a question."
+
+    try:
+        # Define a prompt template
+        template = """You are a helpful assistant answering questions based on the provided context.
+        Use only the information given in the context below to answer the question.
+        If the context doesn't contain the answer, say "The provided context does not contain the answer to this question."
+        Be concise.
+
+        Context:
+        {context}
+
+        Question: {question}
+        Answer:"""
+
+        QA_CHAIN_PROMPT = PromptTemplate(
+            input_variables=["context", "question"],
+            template=template,
+        )
+
+        # Create RetrievalQA chain
+        qa_chain = RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type="stuff",
+            retriever=rag_retriever,
+            chain_type_kwargs={"prompt": QA_CHAIN_PROMPT},
+            return_source_documents=False
+        )
+
+        result = qa_chain.invoke({"query": question})
+        answer = result.get("result", str(result)) if isinstance(result, dict) else str(result)
+
+        return answer.strip()
+
+    except Exception as e:
+        return f"Error answering question: {e}"
+
+def summarize_text(text_to_summarize, max_length=150, min_length=30):
+    """Summarizes the provided text using the LLM."""
+    if llm is None:
+        return "Error: Models not initialized properly."
+
+    if not text_to_summarize or not isinstance(text_to_summarize, str) or len(text_to_summarize.strip()) == 0:
+        return "Error: Please enter text to summarize."
+
+    try:
+        # Create a prompt for summarization
+        prompt = f"Summarize the following text concisely, aiming for {min_length} to {max_length} words:\n\n{text_to_summarize}"
+
+        # Use the pipeline directly for summarization
+        summary_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=max_length,
+            temperature=0.5
+        )
+
+        # Generate summary
+        summary_result = summary_pipeline(prompt, max_length=max_length)[0]['generated_text']
+
+        # Extract the actual summary part
+        summary = summary_result.replace(prompt, '').strip()
+
+        return summary
+
+    except Exception as e:
+        return f"Error summarizing text: {e}"
+
+def draft_text(instructions):
+    """Drafts text based on user instructions using the LLM."""
+    if llm is None:
+        return "Error: Models not initialized properly."
+
+    if not instructions or not isinstance(instructions, str) or len(instructions.strip()) == 0:
+        return "Error: Please enter drafting instructions."
+
+    try:
+        # Drafting prompt
+        prompt = f"Write the following based on these instructions:\n\n{instructions}"
+
+        # Use the pipeline for text generation
+        draft_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=500,
+            temperature=0.7
+        )
+
+        # Generate draft
+        draft_result = draft_pipeline(prompt, max_length=500)[0]['generated_text']
+
+        # Extract the actual draft part
+        draft = draft_result.replace(prompt, '').strip()
+
+        return draft
+
+    except Exception as e:
+        return f"Error drafting text: {e}"
+
+# Gradio Interface
+def create_gradio_interface():
+    with gr.Blocks(theme=gr.themes.Soft()) as iface:
+        gr.Markdown("# Workplace Assistant (GPT-2 Demo)")
+        gr.Markdown("Powered by GPT-2 and Langchain")
+
+        with gr.Tabs():
+            # Document Q&A Tab
+            with gr.TabItem("Document Q&A (RAG)"):
+                gr.Markdown("Load text content from a document, then ask questions about it.")
+                doc_input = gr.Textbox(label="Paste Document Text Here", lines=10, placeholder="Paste the full text content you want to query...")
+                load_button = gr.Button("Process Document")
+                status_output = gr.Textbox(label="Document Status", value=loaded_doc_name, interactive=False)
+                question_input = gr.Textbox(label="Your Question", placeholder="Ask a question about the document...")
+                ask_button = gr.Button("Ask Question")
+                answer_output = gr.Textbox(label="Answer", lines=5, interactive=False)
+
+                load_button.click(
+                    fn=setup_rag_pipeline,
+                    inputs=[doc_input],
+                    outputs=[status_output]
+                )
+                ask_button.click(
+                    fn=answer_question,
+                    inputs=[question_input],
+                    outputs=[answer_output]
+                )
+
+            # Summarization Tab
+            with gr.TabItem("Summarization"):
+                gr.Markdown("Paste text to get a concise summary.")
+                summarize_input = gr.Textbox(label="Text to Summarize", lines=10, placeholder="Paste text here...")
+                summarize_button = gr.Button("Summarize")
+                summary_output = gr.Textbox(label="Summary", lines=5, interactive=False)
+
+                with gr.Accordion("Advanced Options", open=False):
+                    max_len_slider = gr.Slider(minimum=20, maximum=300, value=150, step=10, label="Max Summary Length (approx words)")
+                    min_len_slider = gr.Slider(minimum=10, maximum=100, value=30, step=5, label="Min Summary Length (approx words)")
+
+                summarize_button.click(
+                    fn=summarize_text,
+                    inputs=[summarize_input, max_len_slider, min_len_slider],
+                    outputs=[summary_output]
+                )
+
+            # Drafting Tab
+            with gr.TabItem("Drafting Assistant"):
+                gr.Markdown("Provide instructions for the AI to draft text.")
+                draft_instructions = gr.Textbox(label="Drafting Instructions", lines=5, placeholder="e.g., Draft a short, friendly email to the team.")
+                draft_button = gr.Button("Generate Draft")
+                draft_output = gr.Textbox(label="Generated Draft", lines=10, interactive=False)
+
+                draft_button.click(
+                    fn=draft_text,
+                    inputs=[draft_instructions],
+                    outputs=[draft_output]
+                )
+
+    return iface
+
+# Launch the interface
+if __name__ == "__main__":
+    try:
+        iface = create_gradio_interface()
+        iface.launch(share=True, debug=True)
+    except Exception as e:
+        print(f"Error launching Gradio interface: {e}")