|
import gradio as gr |
|
from transformers import T5Tokenizer, T5ForConditionalGeneration |
|
import torch |
|
from pydantic import BaseModel |
|
import spaces |
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
|
tokenizer = T5Tokenizer.from_pretrained("alpeshsonar/lot-t5-small-filter", legacy=False) |
|
model = T5ForConditionalGeneration.from_pretrained("alpeshsonar/lot-t5-small-filter", torch_dtype=torch.bfloat16).to(device) |
|
|
|
|
|
@spaces.GPU(duration=120) |
|
def generate_text(input_text): |
|
inputs = tokenizer.encode("Extract lots from given text.\n" + input_text, return_tensors="pt").to(device) |
|
outputs = model.generate(inputs, max_new_tokens=1024) |
|
result = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
return result |
|
|
|
iface = gr.Interface(fn=generate_text, inputs="text", outputs="text", title="Line of Therapy") |
|
|
|
|
|
def run(): |
|
|
|
iface.launch(server_name="0.0.0.0", server_port=7860) |
|
|
|
if __name__ == "__main__": |
|
run() |