YAMITEK/quantization_Generator

1import streamlit as st
2from transformers import AutoModelForCausalLM, AutoTokenizer,pipeline
3import torch
4
5st.title("quantization_Generator Fine tunning model")
6
7# Load model and tokenizer
8model_dir = "quantization_model"
9tokenizer = AutoTokenizer.from_pretrained(model_dir)
10model = AutoModelForCausalLM.from_pretrained(model_dir)
11
12
13code_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
14
15
16
17
18inputs_text=st.text_input("Please enter the text",value="def quicksort(arr):")
19
20if st.button("submit"):
21    generated_code = code_generator(inputs_text, max_length=200, num_return_sequences=1)
22    
23    st.write(generated_code[0]["generated_text"])
24