import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import requests

# Initialize the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("PawanKrd/CosmosRP-8k")
model = AutoModelForCausalLM.from_pretrained("PawanKrd/CosmosRP-8k")

st.title("CosmosRP-8k Roleplay Chatbot")

user_input = st.text_input("You:", "")

if st.button("Generate Response"):
    if user_input:
        # Define the endpoint and payload
        endpoint_url = "https://api.pawan.krd/cosmosrp/v1/chat/completions"
        headers = {
            "Content-Type": "application/json"
        }
        data = {
            "model": "cosmosrp",
            "messages": [
                {"role": "system", "content": "You are a roleplay assistant."},
                {"role": "user", "content": user_input}
            ]
        }
        
        # Make the request to the API
        response = requests.post(endpoint_url, headers=headers, json=data)
        result = response.json()
        
        # Extract and display the response
        generated_response = result['choices'][0]['message']['content']
        st.text_area("CosmosRP-8k:", generated_response)
    else:
        st.warning("Please enter some text.")