Replacing inference pipe with manual

app.py

@@ -1,40 +1,25 @@
 import gradio as gr
 import time
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 model_dir = "tinyllama_model"
-#model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
+# Load the quantized model and tokenizer
 model = AutoModelForCausalLM.from_pretrained(model_dir)
 tokenizer = AutoTokenizer.from_pretrained(model_dir)
 
-
-# Load the TinyLlama text generation pipeline
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-#tokenizer = AutoTokenizer.from_pretrained(model_dir)
-
 # Define the inference function
 def generate_text(prompt):
     start_time = time.time()
-    messages = [
-        {
-            "role": "system",
-            "content": "You are a friendly and helpful chatbot",
-        },
-        {"role": "user", "content": prompt},
-    ]
-    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    results = pipe(prompt, max_length=100, num_return_sequences=1)
+    inputs = tokenizer(prompt, return_tensors='pt')
+    # Manually move tensors to quantized int8 if necessary
+    inputs = {key: val.to(torch.int8) if val.dtype == torch.float32 else val for key, val in inputs.items()}
+    outputs = model.generate(**inputs, max_length=100, num_return_sequences=1)
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     end_time = time.time()
     response_time = end_time - start_time
 
-    # Remove system message:
-    assistant_prompt = "<|assistant|>"
-    generated_text = results[0]['generated_text']
-    if assistant_prompt in generated_text:
-        generated_text = generated_text.split(assistant_prompt)[-1].strip()
-    
     return generated_text, f"{response_time:.2f} seconds"
 
 # Create a Gradio interface
@@ -49,4 +34,4 @@ iface = gr.Interface(
 )
 
 # Launch the interface
-iface.launch()
+iface.launch()