import gradio as gr
from PIL import Image
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

#Setting device to cuda
torch.set_default_device("cuda")

# # Ensure GPU usage if available
# device = "cuda" if torch.cuda.is_available() else "cpu"
# torch.set_default_tensor_type('torch.cuda.FloatTensor' if device=='cuda' else 'torch.FloatTensor')

torch.set_default_tensor_type('torch.cuda.FloatTensor')



# Initialize the model and tokenizer
model = AutoModelForCausalLM.from_pretrained("ManishThota/Thota",
                                             torch_dtype=torch.float16, 
                                             device_map="auto",
                                             trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("ManishThota/Thota", trust_remote_code=True)

def predict_answer(question, max_tokens):
    #Set inputs
    text = f" I am Manish Kumar Thota, Works on multimodal large language models and have an interest in machine learning, deep learning, and natural language processing. I can give helpful, detailed, and polite answers to the user's questions. USER:{question}:"
    
    input_ids = tokenizer(text, return_tensors='pt').input_ids.to('cuda')
    
    #Generate the answer
    output_ids = model.generate(
        input_ids,
        max_new_tokens=max_tokens,
        use_cache=True)[0]
    
    return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()

def gradio_predict(question, max_tokens):
    answer = predict_answer(question, max_tokens)
    return answer

# Define the Gradio interface
iface = gr.Interface(
    fn=gradio_predict,
    inputs=[ gr.Textbox(label="Question", placeholder="e.g. Who is Manish Kumar Thota?", scale=4),
            gr.Slider(2, 100, value=25, label="Token Count", info="Choose between 2 and 100")],
    outputs=gr.TextArea(label="Answer"),
    title="Welcome to Manish Kumar Thota's Portfolio",
    description="Ask questions about Manish and get to know him as if you are directly talking to him, this model is still in development so your feedback is highly valuable to improve the model.",
)

# Launch the app
iface.queue().launch(debug=True)