Spaces:
Sleeping
Sleeping
| import os | |
| import re | |
| from datetime import datetime | |
| import gradio as gr | |
| import json | |
| from dotenv import load_dotenv, find_dotenv | |
| _ = load_dotenv(find_dotenv()) | |
| from training.consts import DEFAULT_INPUT_MODEL, SUGGESTED_INPUT_MODELS | |
| from training.trainer import load_training_dataset, load_tokenizer | |
| from training.generate import generate_response, load_model_tokenizer_for_generate | |
| gpu_family = "a100" | |
| model_dir = "model" | |
| model, tokenizer = load_model_tokenizer_for_generate(model_dir) | |
| def get_completion(prompt, model="dolly-v0-70m"): | |
| # Examples from https://www.databricks.com/blog/2023/03/24/hello-dolly-democratizing-magic-chatgpt-open-models.html | |
| instructions = [prompt] | |
| # set some additional pipeline args | |
| pipeline_kwargs = {'torch_dtype': "auto"} | |
| #if gpu_family == "v100": | |
| #pipeline_kwargs['torch_dtype'] = "float16" | |
| #elif gpu_family == "a10" or gpu_family == "a100": | |
| #pipeline_kwargs['torch_dtype'] = "bfloat16" | |
| pipeline_kwargs['max_new_tokens'] = 300 | |
| # Use the model to generate responses for each of the instructions above. | |
| for instruction in instructions: | |
| response = generate_response(instruction, model=model, tokenizer=tokenizer, **pipeline_kwargs) | |
| if response: | |
| print(f"Instruction: {instruction}\n\n{response}\n\n-----------\n") | |
| return response | |
| def greet(input): | |
| prompt = f""" | |
| Text: ```{input}``` | |
| """ | |
| response = get_completion(prompt) | |
| return response | |
| #iface = gr.Interface(fn=greet, inputs="text", outputs="text") | |
| #iface.launch() | |
| #iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Text to find entities", lines=2)], outputs=[gr.HighlightedText(label="Text with entities")], title="NER with dslim/bert-base-NER", description="Find entities using the `dslim/bert-base-NER` model under the hood!", allow_flagging="never", examples=["My name is Andrew and I live in California", "My name is Poli and work at HuggingFace"]) | |
| iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Prompt")], outputs="text") | |
| iface.launch() | |