app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
from threading import Thread

torch.set_num_threads(2)

# Loading the tokenizer and model from Hugging Face's model hub.
tokenizer = AutoTokenizer.from_pretrained("cnmoro/jack-68m-text-structurization")
model = AutoModelForCausalLM.from_pretrained("cnmoro/jack-68m-text-structurization")

# using CUDA for an optimal experience
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# Function to generate model predictions.
def predict(message, history):
    
    model_inputs = tokenizer([
        f"### Structurize: {message}\n\n### Response:\n"
    ], return_tensors="pt").to(device)
    
    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        model_inputs,
        streamer=streamer,
        max_new_tokens=512,
        top_p=0.2,
        top_k=20,
        temperature=0.1,
        repetition_penalty=2.0,
        length_penalty=-0.5,
        num_beams=1,
        prompt_lookup_num_tokens=10
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()  # Starting the generation in a separate thread.
    partial_message = ""
    for new_token in streamer:
        partial_message += new_token
        yield partial_message

# Setting up the Gradio chat interface.
gr.ChatInterface(predict,
                 title="TextStructurization_Jack68m_CPU",
                 description="Pass a text to be structurized"
                 ).launch()  # Launching the web interface.