Create app.py

app.py

@@ -0,0 +1,171 @@
+import gradio as gr
+from datasets import load_dataset, Dataset
+from llama_index.core import PromptTemplate
+from llama_index.core.prompts import ChatMessage
+from llama_index.llms.openai import OpenAI
+from pydantic import BaseModel, Field
+import asyncio
+import numpy as np
+import pandas as pd
+from chromadb import Client
+from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
+import structlog
+
+logger = structlog.get_logger()
+
+logger.info('Loading embedding model')
+embed_fn = SentenceTransformerEmbeddingFunction('BAAI/bge-small-en-v1.5')
+
+
+def load_train_data_and_vectorstore():
+    logger.info("Loading dataset")
+    ds = load_dataset('SetFit/amazon_reviews_multi_en')
+    train_samples_per_class = 50
+    eval_test_samples_per_class = 10
+    train = Dataset.from_pandas(ds['train'].to_pandas().groupby('label').sample(train_samples_per_class, random_state=1234).reset_index(drop=True))
+    reviews = Client().create_collection(
+        name='reviews',
+        embedding_function=embed_fn,
+        get_or_create=True
+    )
+    logger.info("Adding documents to vector store")
+    reviews.add(documents=train['text'], metadatas=[{'rating': x} for x in train['label']], ids=train['id'])
+    return train, reviews
+
+train, reviews = load_train_data_and_vectorstore()
+
+class Rating(BaseModel):
+    rating: int = Field(..., description="Rating of the review", enum=[0, 1, 2, 3, 4])
+
+llm = OpenAI(model="gpt-4o-mini")
+structured_llm = llm.as_structured_llm(Rating)
+
+prompt_tmpl_str = """\
+The review text is below.
+---------------------
+{review}
+---------------------
+Given the review text and not prior knowledge, \
+please attempt to predict the score of the review.
+
+Query: What is the rating of this review?
+Answer: \
+"""
+
+prompt_tmpl = PromptTemplate(
+    prompt_tmpl_str,
+)
+
+async def zero_shot_predict(text):
+    messages = [
+        ChatMessage.from_str(prompt_tmpl.format(review=text))
+    ]
+    response = await structured_llm.achat(messages)
+    return response.raw.rating
+
+few_shot_prompt_tmpl_str = """\
+The review text is below.
+---------------------
+{review}
+---------------------
+Given the review text and not prior knowledge, \
+please attempt to predict the review score of the context. \
+Here are several examples of reviews and their ratings:
+
+{random_few_shot_examples}
+
+Query: What is the rating of this review?
+Answer: \
+"""
+
+few_shot_prompt_tmpl = PromptTemplate(
+    few_shot_prompt_tmpl_str,
+    function_mappings={"random_few_shot_examples": random_few_shot_examples_fn},
+)
+
+rng = np.random.Generator(np.random.PCG64(1234))
+def random_few_shot_examples_fn(**kwargs):
+    if n_samples:=kwargs.get('n_samples'):
+        random_examples = train.shuffle(generator=rng)[:n_samples]
+    else:
+        random_examples = train.shuffle(generator=rng)[:5]
+
+    result_strs = []
+    for text, rating in zip(random_examples['text'], random_examples['label']):
+        result_strs.append(f"Text: {text}\nRating: {rating}")
+    return "\n\n".join(result_strs)
+
+async def random_few_shot_predict(text, n_examples=5):
+    tasks = []
+    for _ in range(3):
+        messages = [
+            ChatMessage.from_str(few_shot_prompt_tmpl.format(review=text, n_samples=n_examples))
+        ]
+        tasks.append(structured_llm.achat(messages, temperature=0.9))
+    results = await asyncio.gather(*tasks)
+    ratings = [r.raw.rating for r in results]
+    # print(ratings)
+    return pd.Series(ratings).mode()[0]
+
+def dynamic_few_shot_examples_fn(**kwargs):
+    n_examples = kwargs.get('n_examples', 5)
+    retrievals = reviews.query(
+        query_texts=[kwargs['review']],
+        n_results=n_examples
+    )
+    result_strs = []
+    documents = retrievals['documents'][0]
+    metadatas = retrievals['metadatas'][0]
+    for document, metadata in zip(documents, metadatas):
+        result_strs.append(f"Text: {document}\nRating: {metadata.get('rating')}")
+    return "\n\n".join(result_strs)
+
+dynamic_few_shot_prompt_tmpl_str = """\
+The review text is below.
+---------------------
+{review}
+---------------------
+Given the review text and not prior knowledge, \
+please attempt to predict the review score of the context. \
+Here are several examples of reviews and their ratings:
+
+{dynamic_few_shot_examples}
+
+Query: What is the rating of this review?
+Answer: \
+"""
+
+dynamic_few_shot_prompt_tmpl = PromptTemplate(
+    dynamic_few_shot_prompt_tmpl_str,
+    function_mappings={"dynamic_few_shot_examples": dynamic_few_shot_examples_fn},
+)
+
+async def dynamic_few_shot_predict(text, n_examples=5):
+    messages = [
+        ChatMessage.from_str(dynamic_few_shot_prompt_tmpl.format(review=text, n_examples=n_examples))
+    ]
+    response = await dynamic_few_shot_structured_llm.achat(messages)
+    return response.raw.rating
+
+def classify(review, num_examples, api_key):
+    llm = OpenAI(model="gpt-4o-mini", api_key=api_key).as_structured_llm(Rating)
+    zero_shot = asyncio.run(zero_shot_predict(review))
+    random_few_shot = asyncio.run(random_few_shot_predict(review, num_examples))
+    dynamic_few_shot = asyncio.run(dynamic_few_shot_predict(review, num_examples))
+    return zero_shot, random_few_shot, dynamic_few_shot
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            api_key = gr.Textbox(label='Openai API Key')
+            n_examples = gr.Slider(minimum=1, maximum=10, value=5, step=1, label='Number of examples to retrieve', interactive=True)
+            review = gr.Textbox(label='Review', interactive=True)
+            submit = gr.Button(value='Submit')
+        with gr.Column():
+            zero_shot_label = gr.Textbox(label='Zero shot', interactive=False)
+            random_few_shot_label = gr.Textbox(label='Random few shot', interactive=False)
+            dynamic_few_shot_label = gr.Textbox(label='Dynamic few shot', interactive=False)
+        
+    submit.click(classify, [review, n_examples], [zero_shot_label, random_few_shot_label, dynamic_few_shot_label])
+
+demo.queue().launch()