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-pje3d-uvelc1"
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]
 if " #" in final_result:
 final_result = final_result.split(" #")[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.