|
import os |
|
|
|
from langchain.memory import ConversationBufferMemory |
|
from langchain.chains import ConversationChain |
|
import langchain.globals |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
import streamlit as st |
|
from langchain_community.llms import HuggingFaceHub |
|
|
|
my_model_id = os.getenv('MODEL_REPO_ID', 'Default Value') |
|
token = os.getenv('HUGGINGFACEHUB_API_TOKEN') |
|
|
|
@st.cache_resource |
|
def load_model(): |
|
|
|
|
|
model = HuggingFaceHub( |
|
repo_id=my_model_id, |
|
task="text-generation", |
|
model_kwargs={ |
|
"max_new_tokens": 512, |
|
"top_k": 30, |
|
"temperature": 0.1, |
|
"repetition_penalty": 1.03, |
|
}, |
|
huggingfacehub_api_token = token, |
|
) |
|
return model |
|
|
|
def demo_miny_memory(model): |
|
|
|
memory = ConversationBufferMemory(llm = model,max_token_limit = 512) |
|
return memory |
|
|
|
def demo_chain(input_text, memory,model): |
|
|
|
llm_conversation = ConversationChain(llm=model,memory=memory,verbose=langchain.globals.get_verbose()) |
|
|
|
chat_reply = llm_conversation.predict(input=input_text) |
|
return chat_reply |