online-lawyer / app.py
Yehor's picture
Update app.py
a4d32a0 verified
import spaces
import torch
print('torch version:', torch.__version__)
import gradio as gr
from unsloth import FastLanguageModel
max_seq_length = 2048
dtype = None
load_in_4bit = True
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "ua-l/gemma-2-9b-legal-steps200-uk", # YOUR MODEL YOU USED FOR TRAINING
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = load_in_4bit,
)
FastLanguageModel.for_inference(model)
def predict(question):
inputs = tokenizer(
[f'''### Question:
{question}
### Answer:
'''], return_tensors = "pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens = 128)
results = tokenizer.batch_decode(outputs, skip_special_tokens=True)
return results[0]
inputs = gr.Textbox(lines=2, label="Enter a question", value="Як отримати виплати ВПО?")
outputs = gr.Markdown(label="Answer")
demo = gr.Interface(fn=predict, inputs=inputs, outputs=outputs)
demo.launch()