import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the fine-tuned model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("ahmed-7124/dgptAW")
model = AutoModelForCausalLM.from_pretrained("ahmed-7124/dgptAW")

# Function to generate response from the model
def doctor_consultant(query):
    # Encode the input query and generate the model's response
    inputs = tokenizer(query, return_tensors="pt")
    outputs = model.generate(inputs['input_ids'], max_length=200, num_return_sequences=1, no_repeat_ngram_size=2, top_p=0.95, temperature=0.7)
    
    # Decode the output and return the response
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("# Doctor Consultant Assistant")
    
    with gr.Row():
        gr.Textbox(label="Ask the Doctor", placeholder="Enter your symptoms or question", lines=3, elem_id="input_text")
    
    with gr.Row():
        gr.Button("Get Response", elem_id="response_button")
    
    with gr.Row():
        gr.Textbox(label="Doctor's Response", elem_id="response_output", interactive=False)
    
    # Connect the function to the interface
    gr.Interface(fn=doctor_consultant, inputs="text", outputs="text").launch(share=True)