Medical_Summary / app.py
adityabhalla-code
updated the name of the application
9c3d2ae
import torch
import gradio
from transformers import AutoModelWithLMHead, AutoTokenizer
def generate_response(model, tokenizer, prompt, max_length=200):
input_ids = tokenizer.encode(prompt, return_tensors="pt") # 'pt' for returning pytorch tensor
# Create the attention mask and pad token id
attention_mask = torch.ones_like(input_ids)
pad_token_id = tokenizer.eos_token_id
output = model.generate(
input_ids,
max_length=max_length,
num_return_sequences=1,
attention_mask=attention_mask,
pad_token_id=pad_token_id
)
return tokenizer.decode(output[0], skip_special_tokens=True)
# Load your model from hub
username = "vsen7" # change it to your HuggingFace username
checkpoint = username + '/Medical_Summary'
loaded_model = AutoModelWithLMHead.from_pretrained(checkpoint)
# Load your tokenizer from hub
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
# Function for response generation
def generate_query_response(prompt):
model = loaded_model
#tokenizer = tokenizer
response = generate_response(model, tokenizer, prompt)
return response
# Gradio interface to generate UI link
iface = gradio.Interface(fn=generate_query_response,
inputs="textbox",
outputs="textbox",
title="Medical Summary",
description="via gradio",
allow_flagging="never",
)
iface.launch()