app.py

from transformers import pipeline
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
from nltk.tokenize import sent_tokenize
import torch

model = "janny127/autotrain-7qmts-cs1er"
tokenizer = AutoTokenizer.from_pretrained(model)

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

def predict(prompt, history):
 # Prompt
 formatted_prompt = (
 f"### Human: {prompt}### Assistant:"
 )
 
 # Generate the Texts
 sequences = pipeline(
 formatted_prompt,
 do_sample=True,
 top_k=50,
 top_p = 0.7,
 num_return_sequences=1,
 repetition_penalty=1.1,
 max_new_tokens=500,
 )
 generated_text = sequences[0]['generated_text']

 final_result = generated_text.split("### Assistant:")[1]
 if " Human: " in final_result:
 final_result = final_result.split(" Human: ")[0]
 
 # return generated_text.strip()
 return final_result.strip()

gr.ChatInterface(predict,
 title="Tinyllama_chatBot",
 description="Ask Tiny llama any questions",
 examples=['How to cook a fish?', 'Who is the president of US now?']
 ).launch() # Launching the web interface.

# interface = gr.ChatInterface(
# fn=predict,
# stop_btn=None
# )

# with gr.Blocks() as demo:
# interface.render()


# demo.launch()

# def generate_answer(query, sample_num=3):
# formatted_prompt = (
# f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"

# )

# sequences = pipeline(
# formatted_prompt,
# do_sample=True,
# top_k=50,
# top_p = 0.9,
# num_return_sequences=sample_num,
# repetition_penalty=1.1,
# max_new_tokens=150,
# eos_token_id=CHAT_EOS_TOKEN_ID,
# )
# answers = list()
# for seq in sequences:
# answer = seq['generated_text'].replace(formatted_prompt, "")
# answers.append(answer)
# answer = sent_tokenize(answers[0])
# final_answer = ''
# for an in answer:
# final_answer += an.strip()
# return final_answer