Spaces:
Build error
Build error
| from langchain.llms.openai import OpenAI | |
| import torch | |
| from PIL import Image, ImageDraw, ImageOps | |
| from transformers import pipeline, BlipProcessor, BlipForConditionalGeneration, BlipForQuestionAnswering | |
| import pdb | |
| class TextRefiner: | |
| def __init__(self, device, api_key=""): | |
| print(f"Initializing TextRefiner to {device}") | |
| self.llm = OpenAI(model_name="gpt-3.5-turbo", temperature=0, openai_api_key=api_key) | |
| self.prompt_tag = { | |
| "imagination": {"True": "could", | |
| "False": "could not"} | |
| } | |
| self.short_prompts = { | |
| "length": "around {length} words", | |
| "sentiment": "of {sentiment} sentiment", | |
| "language": "in {language}", | |
| } | |
| self.long_prompts = { | |
| "imagination": "The new sentence could extend the original description by using your imagination to create additional details, or think about what might have happened before or after the scene in the image, but should not conflict with the original sentence", | |
| } | |
| self.wiki_prompts = "I want you to act as a Wikipedia page. I will give you a sentence and you will parse the single main object in the sentence and provide a summary of that object in the format of a Wikipedia page. Your summary should be informative and factual, covering the most important aspects of the object. Start your summary with an introductory paragraph that gives an overview of the object. The overall length of the response should be around 100 words. You should not describe the parsing process and only provide the final summary. The sentence is \"{query}\"." | |
| self.control_prompts = "As a text reviser, you will convert an image description into a new sentence or long paragraph. The new text is {prompts}. {long_prompts} The sentence is \"{query}\" (give me the revised sentence only)" | |
| def parse(self, response): | |
| out = response.strip() | |
| return out | |
| def parse2(self, response): | |
| out = response.strip() | |
| return out | |
| def prepare_input(self, query, short_prompts, long_prompts): | |
| input = self.control_prompts.format(**{'prompts': ', '.join(short_prompts), 'long_prompts': '. '.join(long_prompts), 'query': query}) | |
| print('prompt: ', input) | |
| return input | |
| def inference(self, query: str, controls: dict, context: list=[]): | |
| """ | |
| query: the caption of the region of interest, generated by captioner | |
| controls: a dict of control singals, e.g., {"length": 5, "sentiment": "positive"} | |
| """ | |
| prompts = [] | |
| long_prompts = [] | |
| for control, value in controls.items(): | |
| # if control in self.prompt_tag: | |
| # value = self.prompt_tag[control][value] | |
| if control in self.short_prompts: | |
| prompts.append(self.short_prompts[control].format(**{control: value})) | |
| else: | |
| if value in [True, "True", "true"]: | |
| long_prompts.append(self.long_prompts[control]) | |
| input = self.prepare_input(query, prompts, long_prompts) | |
| response = self.llm(input) | |
| response = self.parse(response) | |
| tmp_configs = {"query": query} | |
| prompt_wiki = self.wiki_prompts.format(**tmp_configs) | |
| response_wiki = self.llm(prompt_wiki) | |
| response_wiki = self.parse2(response_wiki) | |
| out = { | |
| 'raw_caption': query, | |
| 'caption': response, | |
| 'wiki': response_wiki | |
| } | |
| print(out) | |
| return out | |
| if __name__ == "__main__": | |
| model = TextRefiner(device='cpu') | |
| controls = { | |
| "length": "30", | |
| "sentiment": "negative", | |
| # "imagination": "True", | |
| "imagination": "False", | |
| "language": "English", | |
| } | |
| # model.inference(query='a dog is sitting on a brown bench', controls=controls) | |
| model.inference(query='a cat is sleeping', controls=controls) | |