Update app.py

app.py

@@ -1,31 +1,45 @@
 from operator import ge
 from xml.dom.expatbuilder import theDOMImplementation
 import gradio as gr
-from huggingface_hub import InferenceClient
 import os 
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-def generate(
-    model_name,
-    text,
-    max_new_tokens,
-    top_k
-):    
+# Optional: cache loaded models to avoid reloading every time
+model_cache = {}
+
+def generate(model_name, text, max_new_tokens, top_k):    
     if model_name == "Medium-GPTNeo":
-        model = "tniranjan/finetuned_gptneo-base-tinystories-ta_v3"
+        model_id = "tniranjan/finetuned_gptneo-base-tinystories-ta_v3"
     elif model_name == "Small-GPTNeo":
-        model = "tniranjan/finetuned_tinystories_33M_tinystories_ta"
+        model_id = "tniranjan/finetuned_tinystories_33M_tinystories_ta"
     elif model_name == "Small-LLaMA":
-        model = "tniranjan/finetuned_Llama_tinystories_tinystories_ta"
-    client = InferenceClient(provider="hf-inference",
-    api_key=os.environ["HUGGINGFACEHUB_API_TOKEN"])
+        model_id = "tniranjan/finetuned_Llama_tinystories_tinystories_ta"
+
+    # Load model and tokenizer (from cache if available)
+    if model_id not in model_cache:
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_pretrained(model_id)
+        model_cache[model_id] = (tokenizer, model)
+    else:
+        tokenizer, model = model_cache[model_id]
 
-    return client.text_generation(
-        model=model,
-        prompt = text,
+    inputs = tokenizer(text, return_tensors="pt")
+    
+    # Generate text
+    output = model.generate(
+        **inputs,
         max_new_tokens=max_new_tokens,
         top_k=top_k,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id,
     )
 
+    # Decode generated tokens
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    return generated_text
+    
 demo = gr.Interface(
     generate,
     title="Kurunkathai: Tinystories in Tamil",