app.py

import gradio as gr

from peft import PeftModel, PeftConfig
from transformers import (
    MistralForCausalLM,
    TextIteratorStreamer,
    AutoTokenizer,
    BitsAndBytesConfig,
    GenerationConfig,
)
from time import sleep
from threading import Thread
from torch import float16
import spaces

config = PeftConfig.from_pretrained("lang-uk/dragoman")
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=float16,
    bnb_4bit_use_double_quant=False,
)

model = MistralForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1", quantization_config=quant_config
)


# device_map="auto",)
model = PeftModel.from_pretrained(model, "lang-uk/dragoman").to("cuda")
tokenizer = AutoTokenizer.from_pretrained(
    "mistralai/Mistral-7B-v0.1", use_fast=False, add_bos_token=False
)


@spaces.GPU(duration=30)
def translate(input_text):
    generated_text = ""
    input_text = input_text.strip()

    input_text = f"[INST] {input_text} [/INST]"
    inputs = tokenizer([input_text], return_tensors="pt").to(model.device)

    generation_kwargs = dict(inputs, max_new_tokens=200, num_beams=10, temperature=1) # streamer=streamer,
    
    
    # streaming support
    #streamer = TextIteratorStreamer(
    #    tokenizer, skip_prompt=True, skip_special_tokens=True
    #)

    #thread = Thread(target=model.generate, kwargs=generation_kwargs)

    #thread.start()


    #for new_text in streamer:
    #    generated_text += new_text
    #    yield generated_text

    #generated_text += "\n"
    #yield generated_text

    output = model.generate(**generation_kwargs)
    output = tokenizer.decode(output[0], skip_special_tokens=True).split("[/INST] ")[-1].strip()
    return output


iface = gr.Interface(
    fn=translate,
    inputs=gr.Textbox(
        value="",
        label="Source sentence",
    ),
    outputs=gr.Textbox(label="Translated sentence"),
    examples=[[
        "ChatGPT (Chat Generative Pre-trained Transformer) is a chatbot developed by OpenAI and launched on November 30, 2022. Based on a large language model, it enables users to refine and steer a conversation towards a desired length, format, style, level of detail, and language. Successive prompts and replies, known as prompt engineering, are considered at each conversation stage as a context.[2] ",
        "who holds this neighborhood?"]],
)
iface.launch()