File size: 2,908 Bytes
d3645b8 f35929a d3645b8 |
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 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import pandas as pd
# Load pretrained model and tokenizer
model_name = "jrocha/tiny_llama"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Load data
df = pd.read_csv('splitted_df_jo.csv')
# Function to prepare context
def prepare_context():
pubmed_information_column = df['section_text']
pubmed_information_cleaned = ""
for text in pubmed_information_column.tolist():
objective_index = text.find("Objective")
if objective_index != -1:
cleaned_text = text[:objective_index]
pubmed_information_cleaned += cleaned_text
else:
pubmed_information_cleaned += text
max_length = 1000 # Adjust as needed
return pubmed_information_cleaned[:max_length]
# Function to generate answer
def answer_question(question):
pubmed_information_cleaned = prepare_context()
# Prepare input sequence
messages = [
{
"role": "system",
"content": "You are a friendly chatbot who responds to questions about cancer. Please be considerate.",
},
{"role": "user", "content": question},
]
prompt_with_pubmed = f"{pubmed_information_cleaned}\n\n" # Adjust formatting as needed
prompt_with_pubmed += tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
# Generate response
input_ids = tokenizer.encode(prompt_with_pubmed, return_tensors='pt')
output = model.generate(input_ids, max_length=600, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
# Decode and return generated text
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
position_assistant = generated_text.find("<|assistant|>") + len("<|assistant|>")
return generated_text[position_assistant:]
def main():
""""
Initializes a Women Cancer ChatBot interface using Hugging Face models for question answering.
This function loads a pretrained tokenizer and model from the Hugging Face model hub
and creates a Gradio interface for the ChatBot. Users can input questions related to
women's cancer topics, and the ChatBot will generate answers based on the provided context.
Returns:
None
Example:
>>> main()
"""
iface = gr.Interface(fn=answer_question,
inputs=["text"],
outputs=[gr.Textbox(label="Answer")],
title="Women Cancer ChatBot",
description="How can I help you?",
examples=[
["What is breast cancer?"],
["What are treatments for cervical cancer?"]
])
return iface.launch(debug = True, share=True)
main() |