fix

app.py

@@ -1,10 +1,16 @@
+import os
+
 import streamlit as st
+import torch
+from huggingface_hub import login
+from peft import PeftModel
 from transformers import AutoModelForCausalLM, LlamaTokenizer
 
+login(token=os.getenv("HUGGINGFACE_API_KEY"))
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 @st.cache_resource
 def load():
-    """
     base_model = AutoModelForCausalLM.from_pretrained(
         "stabilityai/japanese-stablelm-instruct-alpha-7b", 
         device_map="auto",
@@ -18,7 +24,6 @@ def load():
         "lora_adapter",
         device_map="auto",
         )
-    """
     model = None
     tokenizer = LlamaTokenizer.from_pretrained(
         "lora_adapter", 
@@ -42,8 +47,26 @@ def get_input_token_length(user_query, system_prompt, messages=""):
     input_ids = tokenizer([prompt], return_tensors='np', add_special_tokens=False)['input_ids']
     return input_ids.shape[-1]
 
-def generate():
-    pass
+def generate_response(user_query: str, system_prompt: str, messages: str="", temperature: float=0, top_k: int=50, top_p: float=0.95, repetition_penalty: float=1.1):
+    prompt = get_prompt(user_query, system_prompt, messages)
+    inputs = tokenizer(
+        prompt, 
+        add_special_tokens=False, 
+        return_tensors="pt"
+        ).to(device)
+    max_new_tokens = 2048 - get_input_token_length(user_query, system_prompt, messages)
+    model.eval()
+    with torch.no_grad():
+        tokens = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens, 
+            temperature=temperature, 
+            top_k=top_k, 
+            top_p=top_p,
+            repetition_penalty=repetition_penalty, 
+            )
+    response = tokenizer.decode(tokens[0][inputs.shape[1]:], skip_special_tokens=True).strip()
+    return response
 
 
 st.header(":dna: 遺伝カウンセリング対話AI")
@@ -90,8 +113,15 @@ if user_prompt := st.chat_input("質問を送信してください"):
     with st.chat_message("user"):
         st.text(user_prompt)
     st.session_state["messages"].append({"role": "user", "content": user_prompt})
-    token_kength = get_input_token_length(user_query=user_prompt, system_prompt=st.session_state["options"]["system_prompt"], messages=st.session_state["messages"])
-    response = f"{token_kength}: " + get_prompt(user_query=user_prompt, system_prompt=st.session_state["options"]["system_prompt"], messages=st.session_state["messages"])
+    response = generate_response(
+        user_prompt=user_prompt, 
+        system_prompt=st.session_state["options"]["system_prompt"],
+        messages=st.session_state["messages"],
+        temperature=st.session_state["options"]["temperature"],
+        top_k=st.session_state["options"]["top_k"],
+        top_p=st.session_state["options"]["top_p"],
+        repetition_penalty=st.session_state["options"]["repetition_penalty"],
+        )
     with st.chat_message("assistant"):
         st.text(response)
     st.session_state["messages"].append({"role": "assistant", "content": response})