Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,245 Bytes
396b713 76a4b71 813d1bb 396b713 813d1bb 6173ae6 396b713 6173ae6 396b713 813d1bb 8a82af7 396b713 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
import spaces
device = 'cuda' if torch.cuda.is_available() else 'cpu'
@spaces.GPU
def load_model():
tokenizer = AutoTokenizer.from_pretrained("alibidaran/medical_transcription_generator")
model = AutoModelForCausalLM.from_pretrained("alibidaran/medical_transcription_generator").to(device)
return model,tokenizer
@spaces.GPU
def generate_text(Text,Max_length,Temperature):
model,tokenizer=load_model()
torch.manual_seed(32)
tokenizer.pad_token_id=tokenizer.eos_token_id
with torch.no_grad():
input_ids = tokenizer(Text, return_tensors="pt")["input_ids"].to(device)
attn_mask=tokenizer(Text, return_tensors="pt")["attention_mask"].to(device)
output=model.generate(input_ids=input_ids,attention_mask=attn_mask,max_new_tokens=Max_length,do_sample=True, temperature=Temperature, top_p=0.90,top_k=10)
return tokenizer.decode(output[0])
demo=gr.Interface(
generate_text,
['text',
gr.Slider(50,2000,value=100,step=10),
gr.Slider(0,2,value=0.7,step=0.1)],
'text',
theme=gr.themes.Base(primary_hue='blue',secondary_hue='cyan'),
description="Medical Trasncript Generator"
)
demo.launch() |