| from transformers import pipeline |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| import gradio as gr |
| from nltk.tokenize import sent_tokenize |
| import torch |
|
|
| model = "janny127/autotrain-5e45b-p5z66" |
| tokenizer = AutoTokenizer.from_pretrained(model) |
|
|
| pipeline = pipeline( |
| "text-generation", |
| model=model, |
| torch_dtype=torch.float16, |
| device_map="auto", |
| ) |
|
|
| def predict(prompt, history): |
| |
| formatted_prompt = ( |
| f"### Human: {prompt}### Assistant:" |
| ) |
| |
| |
| 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 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() |
|
|
| |
