melhora interface e prompt

app.py

@@ -1,35 +1,84 @@
 import os
 
 import gradio as gr
-from generate_text import GenerateStyleText
+from style_scribble import StyleScribble
+
+seed_text = "To serve the intricate and varied demands of image editing, precise and flexible manipulation " \
+            "of image content is indispensable. Recently, DragGAN has achieved impressive editing results" \
+            " through point-based manipulation. However, we have observed that DragGAN struggles with miss " \
+            "tracking, where DragGAN encounters difficulty in effectively tracking the desired handle points, " \
+            "and ambiguous tracking, where the tracked points are situated within other regions that bear " \
+            "resemblance to the handle points. To deal with the above issues,  we propose FreeDrag, which adopts " \
+            "a feature-oriented approach to free the burden on point tracking within the point-oriented methodology" \
+            " of DragGAN. The FreeDrag incorporates adaptive template features, line search, and fuzzy localization " \
+            "techniques to perform stable and efficient point-based image editing. Extensive experiments demonstrate" \
+            " that our method is superior to the DragGAN and enables stable point-based editing in challenging " \
+            "scenarios with similar structures, fine details, or under multi-point targets."
+
+prompt_example = "Write a story about a girl participating in a dancing competition"
 
 
 def process(model, key, repo, example, prompt):
     set_key(key)
-    generate_text = GenerateStyleText(example=example, prompt=prompt)
+    generate_text = StyleScribble(example=example, prompt=prompt)
     if model == "HugginFaceHub":
         model = repo
     generate_text.set_imp_llm(model)
     return generate_text.run()
 
 
-
 def set_key(key):
     os.environ['OPENAI_API_KEY'] = key
     os.environ['HUGGINGFACEHUB_API_TOKEN'] = key
     os.environ['ANTHROPIC_API_KEY'] = key
 
 
-title = "StyleScribble Demo"
 description = "This is a demo of StyleScribble the AI powered app that generates text in your writing style, " \
               "you can learn more and sign-up for full launch here: https://stylescribble.fly.dev"
 
-model_c = gr.Dropdown(choices=["GPT3", "GPT4", "Claude", "HugginFaceHub"], label="Model", value="GPT3")
-key_c = gr.Textbox(label="API Key")
-repo_c = gr.Textbox(label="HF Repo(if selected HugginFaceHub as model)")
-example_c = gr.Textbox(label="Example Writing:", lines=15)
-prompt_c = gr.Textbox(label="Prompt:", lines=5)
-output = gr.Textbox(label="Generated Text:", lines=39)
-demo = gr.Interface(process, [model_c, key_c, repo_c, example_c, prompt_c], output, title=title, description=description)
+with gr.Blocks(css=".no-border {border: none !important;} center: {justify-content: center;}") as demo:
+    with gr.Box(elem_classes="no-border"):
+        with gr.Row():
+            with gr.Box(elem_classes="no-border"):
+                gr.Markdown(
+                    f"""     
+                            # StyleScribble Demo
+                            {description}
+                             """, elem_classes="text")
+                with gr.Row():
+                    with gr.Box(elem_classes="no-border"):
+                        example = gr.Button(value="Example")
+        with gr.Row():
+            with gr.Box(elem_classes="no-border"):
+                model_c = gr.Dropdown(choices=["GPT3", "GPT4", "Claude", "HugginFaceHub"], label="Model",
+                                      value="GPT3")
+                key_c = gr.Textbox(label="API Key")
+                repo_c = gr.Textbox(label="HF Repo", visible=False)
+                example_c = gr.Textbox(label="Paste here a style of text you want to imitate", lines=15)
+                prompt_c = gr.Textbox(label="Write here the desription of the text you want to generate", lines=5)
+            with gr.Box(elem_classes="no-border"):
+                output = gr.Textbox(label="Generated Text:", lines=34, show_copy_button=True)
+        with gr.Row():
+            with gr.Box(elem_classes="no-border center"):
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Row():
+                            gr.ClearButton([key_c, repo_c, example_c, prompt_c, output])
+                            run = gr.Button(variant="primary")
+                    with gr.Column():
+                        pass
+
+
+    def show_example():
+        return [gr.update(value=seed_text), gr.update(value=prompt_example)]
+
+
+    def show_repo_id(model):
+        return gr.update(visible=model == "HugginFaceHub")
+
+
+    run.click(fn=process, inputs=[model_c, key_c, repo_c, example_c, prompt_c], outputs=output)
+    example.click(show_example, [], [example_c, prompt_c])
+    model_c.select(show_repo_id, inputs=[model_c], outputs=[repo_c])
 
 demo.launch()

style_scribble.py

@@ -0,0 +1,80 @@
+from langchain.base_language import BaseLanguageModel
+from langchain.chains import LLMChain, SequentialChain
+from langchain.chat_models import ChatAnthropic
+from langchain.chat_models import ChatOpenAI
+from langchain.llms import HuggingFaceHub
+from langchain.prompts import (
+    PromptTemplate,
+    ChatPromptTemplate,
+    SystemMessagePromptTemplate,
+    HumanMessagePromptTemplate,
+)
+
+
+class StyleScribble:
+    example: str
+    prompt: str
+    llm: BaseLanguageModel
+
+    def __init__(self, example=None, prompt=None, llm=None):
+        self.example = example
+        self.prompt = prompt
+        self.llm = llm
+
+    def set_imp_llm(self, model):
+        if model == 'GPT3':
+            self.llm = ChatOpenAI(model_name="gpt-3.5-turbo-16k")
+        elif model == "GPT4":
+            self.llm = ChatOpenAI(model_name="gpt-4")
+        elif model == "Claude":
+            self.llm = ChatAnthropic()
+        else:
+            self.llm = HuggingFaceHub(repo_id=model)
+
+    def run(self):
+        return self.process()
+
+    def process(self):
+        seq_chain = SequentialChain(
+            chains=[self.get_extract_tone_chain(), self.get_generate_text_chain(self.prompt),
+                    self.get_apply_style_chain()],
+            input_variables=["text"], verbose=True)
+        result = seq_chain({'text': self.example, "style": ""})
+        return str(result.get('result'))
+
+    def create_chain(self, chat_prompt, output_key):
+        return LLMChain(llm=self.llm,
+                        prompt=chat_prompt, output_key=output_key)
+
+    def get_extract_tone_chain(self):
+        template = """Building upon the nuances and distinctive traits in the sample text, establish
+         a comprehensive style guide that encapsulates the unique tone and writing style present in the sample. 
+         This guide should focus on compelling tactics that foster a sense of connection between readers and 
+         the content. Refrain from discussing the specific theme of the sample text or using it as a direct example. 
+         Instead, formulate your analysis in such a way that it remains abstract, facilitating its application to any
+         other text that might be inspired by or originate from the sample. This abstract analysis will enable writers to adopt the essence of the style while allowing for versatility across various themes or topics..
+        """
+        system_message_prompt = SystemMessagePromptTemplate.from_template(template)
+        human_template = "{text}"
+        human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
+        chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
+
+        return self.create_chain(chat_prompt, "style")
+
+    def get_generate_text_chain(self, prompt):
+        template = """Generate a text following the user_request(use same language of the request):
+                      {user_request}
+                      """.replace("{user_request}", prompt)
+        return self.create_chain(PromptTemplate.from_template(template),
+                                 "generated_text")
+
+    def get_apply_style_chain(self):
+        template = """STYLE:
+                    {style}
+                    REWRITE THE TEXT BELLOW APPLYING THE STYLE GUIDE ABOVE(use same language of the request):
+                    {generated_text}
+        """
+
+        prompt = PromptTemplate.from_template(template=template)
+        prompt.partial(style="")
+        return self.create_chain(prompt, "result")