|
import gradio as gr |
|
from transformers import pipeline |
|
import logging |
|
import re |
|
from llama_cpp import Llama |
|
from huggingface_hub import hf_hub_download |
|
import sys |
|
|
|
|
|
logging.basicConfig(level=logging.INFO) |
|
|
|
|
|
model_repo="filipealmeida/open-llama-3b-v2-pii-transform" |
|
model_filename="ggml-model-f16.gguf" |
|
|
|
def download_model(): |
|
print("Downloading model...") |
|
sys.stdout.flush() |
|
file = hf_hub_download( |
|
repo_id=model_repo, filename=model_filename |
|
) |
|
print("Downloaded " + file) |
|
return file |
|
|
|
def generate_text(prompt, example): |
|
logging.debug(f"Received prompt: {prompt}") |
|
input = f""" |
|
### Instruction: |
|
{prompt} |
|
### Response: |
|
""" |
|
|
|
logging.info(f"Input : {input}") |
|
|
|
output = llm(input, max_tokens=200, stop=["\n"]) |
|
print(output) |
|
generated_text = output['choices'][0]['text'] |
|
logging.info(f"Generated text: {generated_text}") |
|
|
|
parsed_text = generated_text.split("\n")[0] |
|
|
|
logging.info(f"Parsed text: {parsed_text}") |
|
return parsed_text |
|
|
|
model = download_model() |
|
llm = Llama(model_path=model) |
|
|
|
|
|
interface = gr.Interface( |
|
fn=generate_text, |
|
inputs=[ |
|
gr.Textbox(lines=1, placeholder="Enter text to anonimize...", label="Prompt", |
|
value="My name is Filipe and my phone number is 555-121-2234. How are you?") |
|
], |
|
outputs=gr.Textbox(label="Generated text") |
|
) |
|
|
|
|
|
interface.launch() |
|
|