Update app.py

app.py

@@ -1,169 +1,100 @@
-from dotenv import load_dotenv
 import streamlit as st
 from langchain_community.document_loaders import UnstructuredPDFLoader
-from langchain_text_splitters.character import CharacterTextSplitter
-from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_groq import ChatGroq
-from langchain.memory import ConversationBufferMemory
-from langchain.chains import ConversationalRetrievalChain
-from nltk import word_tokenize, FreqDist, sent_tokenize
-from textblob import TextBlob
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.decomposition import NMF
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, pipeline
+import torch
 import os
-import nltk
-nltk.download('punkt')
-
-# Load secret API key
-secret = os.getenv('Groq_api')
-working_dir = os.path.dirname(os.path.abspath(__file__))
-
-def load_documents(file_path):
-    loader = UnstructuredPDFLoader(file_path)
-    documents = loader.load()
-    return documents
-
-def setup_vectorstore(documents):
-    embeddings = HuggingFaceEmbeddings()
-    text_splitter = CharacterTextSplitter(
-        separator="\n",
-        chunk_size=500,
-        chunk_overlap=100
-    )
-    doc_chunks = text_splitter.split_documents(documents)
-    vectorstores = FAISS.from_documents(doc_chunks, embeddings)
-    return vectorstores
-
-def extract_mentality_traits(documents):
-    # Extract the text from each document
-    text = "\n".join([doc.page_content for doc in documents])
-    blob = TextBlob(text)
-
-    # Analyze tone
-    sentiment = blob.sentiment.polarity
-    tone = "neutral"
-    if sentiment > 0.2:
-        tone = "positive"
-    elif sentiment < -0.2:
-        tone = "negative"
-
-    # Common phrases and reasoning patterns
-    words = word_tokenize(text.lower())
-    fdist = FreqDist(words)
-    common_phrases = [word for word, count in fdist.most_common(10) if len(word) > 3]
-
-    # Topic Modeling to identify key viewpoints
-    vectorizer = TfidfVectorizer(stop_words='english')
-    doc_term_matrix = vectorizer.fit_transform([text])
-    nmf = NMF(n_components=2, random_state=1)
-    nmf.fit(doc_term_matrix)
-    topics = []
-    for topic in nmf.components_:
-        topic_words = [vectorizer.get_feature_names_out()[i] for i in topic.argsort()[-5:]]
-        topics.append(", ".join(topic_words))
-
-    mentality_traits = {
-        "tone": tone,
-        "common_phrases": common_phrases,
-        "reasoning_pattern": "deductive" if sentiment > 0 else "inductive",
-        "key_viewpoints": topics
-    }
-    return mentality_traits
-
-from langchain.prompts import PromptTemplate
-
-from langchain.prompts import PromptTemplate
-from langchain.chains import LLMChain
-
-from langchain.prompts import PromptTemplate
-from langchain.chains import LLMChain
-
-def create_chain(vectorstores, mentality_traits):
-    # Create style prompt dynamically based on mentality traits
-    style_prompt = (
-        f"Think and respond in the same mindset as the person from the uploaded document. "
-        f"The author's tone is generally {mentality_traits['tone']}. "
-        f"Their reasoning style is mostly {mentality_traits['reasoning_pattern']} and they often discuss themes such as {', '.join(mentality_traits['key_viewpoints'])}. "
-        f"Use phrases like {', '.join(mentality_traits['common_phrases'])}. "
-        f"Respond as if you are the author, mirroring their values, reasoning style, and mentality in all replies."
-    )
-
-    # Create a custom prompt template for conversational retrieval
-    prompt = PromptTemplate(input_variables=["question"], template=style_prompt + "\nQuestion: {question}\nAnswer:")
-
-    llm = ChatGroq(
-        api_key=secret,
-        model="llama-3.1-70b-versatile",
-        temperature=0.7
-    )
-
-    # Create the LLM chain that includes the custom style prompt
-    llm_chain = LLMChain(llm=llm, prompt=prompt)
-
-    retriever = vectorstores.as_retriever()
-    memory = ConversationBufferMemory(
-        llm=llm,
-        output_key="answer",
-        memory_key="chat_history",
-        return_messages=True
-    )
-
-    # Now, create the ConversationalRetrievalChain without the `prompt`
-    chain = ConversationalRetrievalChain.from_llm(
-        llm=llm_chain,  # Use the LLMChain here, not the LLM directly
-        retriever=retriever,
-        memory=memory,
-        verbose=True
-    )
-
-    return chain
-
-# Streamlit setup
+from datasets import Dataset
+import pandas as pd
+import re
+
+# Set up page
 st.set_page_config(
-    page_title="Chat with your documents",
-    page_icon="📑",
+    page_title="Tweet Style Cloning",
+    page_icon="🐦",
     layout="centered"
 )
+st.title("🐦 Clone Tweet Style from PDF")
 
-st.title("📝Chat With Your Docs 😎")
-
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = []
+# Step 1: Upload PDF
+uploaded_file = st.file_uploader("Upload a PDF with tweets")
 
-uploaded_file = st.file_uploader(label="Upload your PDF")
+if uploaded_file is not None:
+    # Step 2: Extract text from PDF
+    def load_pdf_text(file_path):
+        loader = UnstructuredPDFLoader(file_path)
+        documents = loader.load()
+        return " ".join([doc.page_content for doc in documents])
 
-if uploaded_file:
-    file_path = f"{working_dir}/{uploaded_file.name}"
-    with open(file_path, "wb") as f:
+    # Save the uploaded PDF file temporarily
+    with open("uploaded_tweets.pdf", "wb") as f:
         f.write(uploaded_file.getbuffer())
 
-    # Load and process the PDF document
-    documents = load_documents(file_path)
-    if "vectorstores" not in st.session_state:
-        st.session_state.vectorstores = setup_vectorstore(documents)
-
-    # Extract mentality traits and create the conversational chain with mental adaptation
-    mentality_traits = extract_mentality_traits(documents)
-    if "conversation_chain" not in st.session_state:
-        st.session_state.conversation_chain = create_chain(st.session_state.vectorstores, mentality_traits)
-
-for message in st.session_state.chat_history:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-user_input = st.chat_input("Ask any questions relevant to uploaded PDF")
-
-if user_input:
-    st.session_state.chat_history.append({"role": "user", "content": user_input})
-    with st.chat_message("user"):
-        st.markdown(user_input)
-
-    # Use the conversation chain to generate the response
-    response = st.session_state.conversation_chain({"question": user_input})
-    
-    # Ensure that assistant_response is cast as a string
-    assistant_response = str(response.get("answer", "Sorry, I couldn't generate a response."))
-    st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
-    with st.chat_message("assistant"):
-        st.markdown(assistant_response)
+    # Extract text from PDF
+    extracted_text = load_pdf_text("uploaded_tweets.pdf")
+
+    # Step 3: Preprocess text to separate each tweet (assuming tweets end with newline)
+    tweets = re.split(r'\n+', extracted_text)
+    tweets = [tweet.strip() for tweet in tweets if len(tweet.strip()) > 0]
+
+    # Display a few sample tweets for verification
+    st.write("Sample Tweets Extracted:")
+    st.write(tweets[:5])
+
+    # Step 4: Fine-tune a model on the extracted tweets
+    def fine_tune_model(tweets):
+        # Convert tweets to a DataFrame and Dataset
+        df = pd.DataFrame(tweets, columns=["text"])
+        tweet_dataset = Dataset.from_pandas(df)
+
+        # Load model and tokenizer
+        model_name = "gpt2"  # Replace with a suitable model if needed
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+
+        # Tokenize the dataset
+        def tokenize_function(examples):
+            return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)
+
+        tokenized_tweets = tweet_dataset.map(tokenize_function, batched=True)
+
+        # Training arguments
+        training_args = TrainingArguments(
+            output_dir="./fine_tuned_tweet_model",
+            per_device_train_batch_size=4,
+            num_train_epochs=3,
+            save_steps=10_000,
+            save_total_limit=1,
+            logging_dir='./logs',
+        )
+
+        # Initialize the Trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_tweets,
+        )
+
+        # Fine-tune the model
+        trainer.train()
+
+        # Save the fine-tuned model
+        model.save_pretrained("fine_tuned_tweet_model")
+        tokenizer.save_pretrained("fine_tuned_tweet_model")
+
+        return model, tokenizer
+
+    # Trigger fine-tuning and notify user
+    with st.spinner("Fine-tuning model..."):
+        model, tokenizer = fine_tune_model(tweets)
+    st.success("Model fine-tuned successfully!")
+
+    # Step 5: Set up text generation
+    tweet_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+    # Generate a new tweet based on user input
+    prompt = st.text_input("Enter a prompt for a new tweet in the same style:")
+    if prompt:
+        with st.spinner("Generating tweet..."):
+            generated_tweet = tweet_generator(prompt, max_length=50, num_return_sequences=1)
+            st.write("Generated Tweet:")
+            st.write(generated_tweet[0]["generated_text"])