from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig
import transformers
import torch


model_id = "QuantumIntelligence/QI-neural-chat-7B-ko-DPO" 

tokenizer = AutoTokenizer.from_pretrained(model_id)
# model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_8bit=True) # quantization

pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    torch_dtype=torch.float16,
    device_map="auto",
    tokenizer=tokenizer,
)

prompt = """Classify the text into neutral, negative or positive. 
Text: This movie is definitely one of my favorite movies of its kind. The interaction between respectable and morally strong characters is an ode to chivalry and the honor code amongst thieves and policemen.
Sentiment:
"""

outputs = pipeline(prompt, max_new_tokens=6)
print(outputs[0]["generated_text"])

prompt = """
다음 텍스트를 중립, 부정, 긍정으로 분류해줘.
텍스트: 하늘을 보니 비가 올듯 하다. 우울한 기분이 들어서 술을 한잔 할까 고민중인데 같이 마실 사람이 없다.
분류:
"""

outputs = pipeline(prompt, max_new_tokens=6)
print(outputs[0]["generated_text"])

prompt = """
국내 연구진이 미국, 영국 공동 연구팀과 청각 기능에 관여하는 단백질 구조를 규명했다. 난청 치료법을 개발하는 데 도움이 될 것으로 보인다.
포스텍은 조윤제 생명과학과 교수 연구팀이 김광표 경희대 응용화학과 교수 연구팀, 브셰볼로드 카트리치 미국 서던 캘리포니아대 교수 연구팀, 캐롤 로빈슨 영국 옥스퍼드대 교수와 함께 청각 관련 특정 수용체 단백질 구조와 메커니즘을 밝히는 데 성공했다고 11일 밝혔다.
귀 안쪽에는 소리를 감지하는 달팽이관과 평형감각을 담당하는 전정기관이 있다. 이 기관들의 세포들은 수용체 단백질인 ‘GPR156’을 갖고 있다. GPR156이 활성화되면 세포 내 G단백질과 결합해 신호를 전달한다. G단백질은 ‘구아닌 뉴클레오타이드-결합 단백질’로 신호를 전달하는 중개자다.
GPR156은 다른 수용체와 달리 특별한 자극이 없어도 항상 높은 활성을 유지하며 청각과 평형 기능 유지에 큰 역할을 한다. 선천적으로 청각 장애가 있는 환자들을 치료하기 위해서는 이 단백질의 구조와 작용 메커니즘을 알아야 한다.
연구팀은 초저온전자현미경(Cryo-EM) 분석법을 사용해 GPR156과 GPR156-G단백질 결합 복합체를 고해상도로 관찰했다. 이를 통해 수용체를 활성화하는 작용제 없이도 GPR156이 높은 활성을 유지할 수 있는 원인을 찾았다.
GPR156은 세포막에 풍부한 인지질과 결합해 활성화됐다. 세포질에 있는 G단백질과의 상호작용을 통해 자체적으로 구조를 변형, 높은 활성을 유지한다는 사실도 확인됐다.
기존에 알려진 수용체 단백질들과 달리 GPR156은 세포막을 통과하는 7번째 힐릭스 말단 부분의 구조를 유연하게 바꾸며 G단백질과의 결합을 유도했다. 이를 통해 신호를 활성화함으로써 소리를 감지하는 데 도움을 주었다.
조 교수는 “선천적으로 난청과 균형 감각 기능에 장애가 있는 환자들이 많다”며 “이들을 위한 획기적인 치료법과 약물 개발에 이번 연구가 큰 도움이 되길 바란다”고 말했다. 연구 논문은 국제학술지 ‘네이처 구조&분자 생물학’ 온라인판에 최근 게재됐다. 

위 문장을 한글로 100자내로 요약해줘.
요약:
"""

outputs = pipeline(prompt, max_new_tokens=256, return_full_text = False, pad_token_id=tokenizer.eos_token_id)&&
print(outputs[0]["generated_text"])

prompt = """
참가자들은 먼저 fMRI 기기 안에서 자신의 이야기를 읽는 동안 뇌의 활동 패턴을 기록했다. 이야기를 다시 읽으면서는 이야기 속 단어에 대해 순간순간 자신이 느끼는 자기 관련도, 긍·부정 정서를 보고했다. 수집된 49명의 데이터는 자기 관련도와 긍·부정 정서 점수에 따라 다섯 개 수준으로 분류됐다.
질문: 실험의 대상이 된 사람은 몇 명인가? 한글로 대답.
대답:
"""

outputs = pipeline(prompt, max_new_tokens=30, return_full_text = False)
generated_text = outputs[0]["generated_text"]
print(generated_text)

prompt = """
각 방에 공이 5개 있고, 방의 총 개수는 4. 총 공의 갯수는 몇개 인가?
"""

outputs = pipeline(prompt, max_new_tokens=40, return_full_text = False, pad_token_id=tokenizer.eos_token_id)
print(outputs[0]["generated_text"])

messages = [{"role": "user", "content": "좋은 취미를 가지려면 어떻게 하나요?"}]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

outputs = pipeline(prompt, max_new_tokens=512, do_sample=True, temperature=0.7, top_k=50, top_p=0.95, return_full_text = False) 
generated_text = outputs[0]["generated_text"]

print(generated_text)