import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from sentence_transformers import SentenceTransformer
from langchain.vectorstores import Chroma
import os

# Hugging Face 모델 ID
model_id = "hewoo/hehehehe"
token = os.getenv("HF_API_TOKEN")  # 필요한 경우 사용자에게 Hugging Face API 토큰 입력을 요청할 수 있음

# 모델과 토크나이저 로드
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)

# 텍스트 생성 파이프라인 설정
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.5, top_p=0.85, top_k=40, repetition_penalty=1.2)

# 사용자 정의 임베딩 클래스 생성
class CustomEmbedding:
    def __init__(self, model):
        self.model = model

    def embed_query(self, text):
        return self.model.encode(text, convert_to_tensor=True).tolist()

    def embed_documents(self, texts):
        return [self.model.encode(text, convert_to_tensor=True).tolist() for text in texts]

# 임베딩 모델 및 벡터 스토어 설정
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
embedding_function = CustomEmbedding(embedding_model)

# Chroma 벡터 스토어 설정
persist_directory = "./chroma_batch_vectors"  # Spaces 환경에 맞게 조정 필요
vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embedding_function)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# 질문에 대한 응답 생성 함수
def generate_response(user_input):
    search_results = retriever.get_relevant_documents(user_input)
    context = "\n".join([result.page_content for result in search_results])
    input_text = f"맥락: {context}\n질문: {user_input}"
    response = pipe(input_text)[0]["generated_text"]
    return response

# Streamlit 앱 UI
st.title("챗봇 test")
st.write("Llama 3.2-3B 모델을 사용한 챗봇입니다. 질문을 입력해 주세요.")

# 사용자 입력 받기
user_input = st.text_input("질문")
if user_input:
    response = generate_response(user_input)
    st.write("챗봇 응답:", response)