MoinRomanticbot / app.py
syedmoinms's picture
Update app.py
a1bbc6a verified
raw
history blame
1.07 kB
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from memory import update_memory, check_memory
with open("persona.txt", "r", encoding="utf-8") as f:
personality = f.read()
model_name = "TheBloke/Pygmalion-7B-GPTQ"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", trust_remote_code=True)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
def chatbot(input_text):
memory_response = check_memory(input_text)
if memory_response:
return memory_response
prompt = f"{personality}\nUser: {input_text}\nAI:"
outputs = generator(prompt, max_length=200, do_sample=True, temperature=0.7, top_p=0.9)
response = outputs[0]["generated_text"].split("AI:")[-1].strip()
update_memory(input_text, response)
return response
iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", title="MoinRomanticBot")
if __name__ == "__main__":
iface.launch(server_name="0.0.0.0", server_port=7860)