Create app.py

app.py

@@ -0,0 +1,176 @@
+import os
+import openai
+from openai import OpenAI
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from datasets import load_dataset, Dataset
+from sklearn.neighbors import NearestNeighbors
+import numpy as np
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, TextStreamer
+import torch
+from typing import List
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+import gradio as gr
+
+# Configuration
+
+# Sample questions:
+# 1. What are the key components of an effective persona for prompt generation?
+# 2. How can I create a persona that generates creative writing prompts?
+# 3. What strategies can I use to make my persona-driven prompts more engaging?
+
+DEFAULT_QUESTION = "Ask me anything in the context of persona-driven prompt generation..."
+
+# Set API keys (make sure these are set in your environment)
+os.environ['OPENAI_BASE'] = "https://api.openai.com/v1"
+os.environ['OPENAI_MODEL'] = "gpt-4"
+os.environ['MODEL_PROVIDER'] = "huggingface"
+api_key = os.environ.get("OPENAI_API_KEY")
+model_provider = os.environ.get("MODEL_PROVIDER")
+
+# Instantiate the client for openai v1.x
+if model_provider.lower() == "openai":
+    MODEL_NAME = os.environ['OPENAI_MODEL']
+    client = OpenAI(
+        base_url=os.environ.get("OPENAI_BASE"),
+        api_key=api_key
+    )
+else:
+    MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"
+    # Initialize Hugging Face client with streaming support
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=os.environ.get("HF_TOKEN"))
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map='auto',
+        token=os.environ.get("HF_TOKEN"),
+        torch_dtype=torch.bfloat16,
+    )
+    streamer = TextStreamer(tokenizer, skip_prompt=True)
+    question_answerer = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device_map='auto',
+        streamer=streamer,
+        max_new_tokens=512,
+        return_full_text=False
+    )
+
+# Load the Hugging Face dataset
+dataset = load_dataset('tosin2013/persona-driven-prompt-generator', streaming=True)
+dataset = Dataset.from_list(list(dataset['train']))
+
+# Initialize embeddings
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# Extract texts from the dataset
+texts = dataset['input']
+
+# Create embeddings for the texts
+text_embeddings = embeddings.embed_documents(texts)
+
+# Fit a nearest neighbor model
+nn = NearestNeighbors(n_neighbors=5, metric='cosine')
+nn.fit(np.array(text_embeddings))
+
+def get_relevant_documents(query, k=5):
+    """
+    Retrieves the k most relevant documents to the query.
+    """
+    query_embedding = embeddings.embed_query(query)
+    distances, indices = nn.kneighbors([query_embedding], n_neighbors=k)
+    relevant_docs = [texts[i] for i in indices[0]]
+    return relevant_docs
+
+def generate_response(question, history):
+    try:
+        print(f"\n[LOG] Received question: {question}")
+        
+        # Get relevant documents based on the query
+        relevant_docs = get_relevant_documents(question, k=3)
+        print(f"[LOG] Retrieved {len(relevant_docs)} relevant documents")
+        
+        # Create the prompt for the LLM
+        context = "\n".join(relevant_docs)
+        prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"
+        print(f"[LOG] Generated prompt: {prompt[:200]}...")  # Log first 200 chars of prompt
+        
+        if model_provider.lower() == "huggingface":
+            prompt_template = """
+            <s>[INST] <<SYS>>
+            You are a helpful AI assistant. Answer the question based on the provided context.
+            <</SYS>>
+        
+            {prompt}[/INST]
+            """
+            chat_prompt = ChatPromptTemplate.from_template(prompt_template)
+            result = question_answerer(chat_prompt.format(prompt=prompt))
+            response = result[0]['generated_text'] if isinstance(result, list) else result
+            print(f"[LOG] Using Hugging Face model: {MODEL_NAME}")
+            print(f"[LOG] Hugging Face response: {response[:200]}...")  # Log first 200 chars of response
+        elif model_provider.lower() == "openai":
+            response = client.chat.completions.create(
+                model=os.environ.get("OPENAI_MODEL"),
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant. Answer the question based on the provided context."},
+                    {"role": "user", "content": prompt},
+                ]
+            )
+            response = response.choices[0].message.content
+            print(f"[LOG] Using OpenAI model: {os.environ.get('OPENAI_MODEL')}")
+            print(f"[LOG] OpenAI response: {response[:200]}...")  # Log first 200 chars of response
+        
+        # Update chat history with new message pair
+        history.append((question, response))
+        return history
+    except Exception as e:
+        error_msg = f"Error generating response: {str(e)}"
+        print(f"[ERROR] {error_msg}")
+        history.append((question, error_msg))
+        return history
+
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown(f"""
+    ## Persona-Driven Prompt Generator QA Agent
+    **Current Model:** {MODEL_NAME}
+    
+    The Custom Prompt Generator is a Python application that leverages Large Language Models (LLMs) and the LiteLLM library to dynamically generate personas, fetch knowledge sources, resolve conflicts, and produce tailored prompts. This application is designed to assist in various software development tasks by providing context-aware prompts based on user input and predefined personas.
+
+    Sample questions:
+    1. What are the key components of an effective persona for prompt generation?
+    2. How can I create a persona that generates creative writing prompts?
+    3. What are the main features of the persona generator?
+    
+    Related repository: [persona-driven-prompt-generator](https://github.com/tosin2013/persona-driven-prompt-generator)
+    """)
+    
+    with gr.Row():
+        chatbot = gr.Chatbot(label="Chat History")
+    
+    with gr.Row():
+        question = gr.Textbox(
+            value=DEFAULT_QUESTION,
+            label="Your Question",
+            placeholder=DEFAULT_QUESTION
+        )
+    
+    with gr.Row():
+        submit_btn = gr.Button("Submit")
+        clear_btn = gr.Button("Clear")
+    
+    # Event handlers
+    submit_btn.click(
+        generate_response,
+        inputs=[question, chatbot],
+        outputs=[chatbot]
+    )
+    
+    clear_btn.click(
+        lambda: (None, ""),
+        inputs=[],
+        outputs=[chatbot, question]
+    )
+
+if __name__ == "__main__":
+    demo.launch()