import gradio as gr
from transformers import AutoTokenizer, AutoModel
import openai
import os

# Load the NASA-specific bi-encoder model and tokenizer
bi_encoder_model_name = "nasa-impact/nasa-smd-ibm-st-v2"
bi_tokenizer = AutoTokenizer.from_pretrained(bi_encoder_model_name)
bi_model = AutoModel.from_pretrained(bi_encoder_model_name)

# Set up OpenAI API key
openai.api_key = os.getenv('OPENAI_API_KEY')

def encode_text(text):
    inputs = bi_tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
    outputs = bi_model(**inputs)
    # Ensure the output is 2D by averaging the last hidden state along the sequence dimension
    return outputs.last_hidden_state.mean(dim=1).detach().numpy().flatten()

def generate_response(user_input, context_embedding):
    # Create a structured prompt for GPT-4
    context_str = ' '.join(map(str, context_embedding))  # Convert context embedding to a string
    combined_input = f"Question: {user_input}\nContext: {context_str}\nAnswer:"
    
    # Generate a response using GPT-4
    response = openai.ChatCompletion.create(
        model="gpt-4",  # Use GPT-4 model
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": combined_input}
        ],
        max_tokens=150,
        temperature=0.5,
        top_p=0.9,
        frequency_penalty=0.5,
        presence_penalty=0.0
    )
    return response.choices[0].message['content'].strip()

def chatbot(user_input, context=""):
    context_embedding = encode_text(context) if context else ""
    response = generate_response(user_input, context_embedding)
    return response

# Create the Gradio interface
iface = gr.Interface(
    fn=chatbot,
    inputs=[gr.Textbox(lines=2, placeholder="Enter your message here..."), gr.Textbox(lines=2, placeholder="Enter context here (optional)...")],
    outputs="text",
    title="Context-Aware Dynamic Response Chatbot",
    description="A chatbot using a NASA-specific bi-encoder model to understand the input context and GPT-4 to generate dynamic responses."
)

# Launch the interface
iface.launch()