Update app.py

app.py

@@ -1,35 +1,23 @@
 import gradio as gr
-import requests
-import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Get the Hugging Face API token from environment variables
-API_URL = "https://api-inference.huggingface.co/models/EmTpro01/codellama-Code-Generator"
-api_token = os.getenv("HF_API_TOKEN")
-headers = {"Authorization": f"Bearer {api_token}"}
+# Load your fine-tuned model and tokenizer
+model_name = "EmTpro01/codellama-Code-Generator"  # Use your model name here
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Function to generate code from a prompt using the Inference API
+# Function to generate code from a prompt
 def generate_code(prompt):
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_length": 150,
-            "temperature": 0.7,
-            "top_k": 50
-        }
-    }
-    response = requests.post(API_URL, headers=headers, json=payload)
-    if response.status_code == 200:
-        generated_code = response.json()[0]['generated_text']
-        return generated_code
-    else:
-        return f"Error: {response.status_code}, {response.text}"
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(inputs.input_ids, max_length=150, temperature=0.7, top_k=50)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Create the Gradio interface
 interface = gr.Interface(
     fn=generate_code, 
     inputs="text", 
     outputs="text", 
-    title="Code Generator using Inference API",
+    title="Code Generator",
     description="Enter a code prompt to generate Python code using the fine-tuned model."
 )