Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import gradio as gr | |
| import spaces | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| def load_model(): | |
| tokenizer = AutoTokenizer.from_pretrained("alibidaran/medical_transcription_generator") | |
| model = AutoModelForCausalLM.from_pretrained("alibidaran/medical_transcription_generator").to(device) | |
| return model,tokenizer | |
| 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() |