Chegue100 commited on
Commit
77f7f04
·
verified ·
1 Parent(s): e23df81

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -5,14 +5,14 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
5
  tokenizer = AutoTokenizer.from_pretrained("sambanovasystems/SambaLingo-Hungarian-Chat", use_fast=False)
6
  model = AutoModelForCausalLM.from_pretrained("sambanovasystems/SambaLingo-Hungarian-Chat", device_map="auto", torch_dtype="auto")
7
 
8
- # Create the pipeline with max_new_tokens parameter
9
- pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto", use_fast=False, max_new_tokens=50)
10
 
11
  # Define the chat function
12
  def chat(question):
13
  messages = [{"role": "user", "content": question}]
14
  prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, truncation=True)
15
- outputs = pipe(prompt)[0]
16
  return outputs["generated_text"]
17
 
18
  # Set up the Gradio interface
 
5
  tokenizer = AutoTokenizer.from_pretrained("sambanovasystems/SambaLingo-Hungarian-Chat", use_fast=False)
6
  model = AutoModelForCausalLM.from_pretrained("sambanovasystems/SambaLingo-Hungarian-Chat", device_map="auto", torch_dtype="auto")
7
 
8
+ # Create the pipeline
9
+ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto", use_fast=False)
10
 
11
  # Define the chat function
12
  def chat(question):
13
  messages = [{"role": "user", "content": question}]
14
  prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, truncation=True)
15
+ outputs = pipe(prompt, max_length=50, truncation=True)[0]
16
  return outputs["generated_text"]
17
 
18
  # Set up the Gradio interface